diff --git a/.gitignore b/.gitignore index e09036933..f8bbccf02 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,21 @@ !.env.example !backend/cmd/installer/files/links/.env +# local build artifacts (custom binary / image build context) +.build_ctx/ +backend/.build_ctx/ +backend/pentagi-kbsearch +# local LLM provider config used by the compose mount +agnes.provider.yml +agnes.provider.yml.bak_* +# installer link stubs generated by the setup wizard (not source) +example.*.provider.yml +# local verification scripts and their logs +rag_closure_*.py +verify_registration.py +kb_test_*.log +rag_kbsearch_run.log + backend/tmp backend/build backend/data diff --git a/Dockerfile.ollama b/Dockerfile.ollama new file mode 100644 index 000000000..4030d23e7 --- /dev/null +++ b/Dockerfile.ollama @@ -0,0 +1,15 @@ +# 本地 Ollama（musl 原生，供 PentAGI 完整 RAG embedding 用） +# 基于 alpine + 清华 apk 镜像源；官方 ollama/ollama 镜像大层拉取反复被重置，故自建 +FROM alpine:3.23 + +RUN sed -i 's#https://dl-cdn.alpinelinux.org#https://mirrors.tuna.tsinghua.edu.cn#g; s#http://dl-cdn.alpinelinux.org#https://mirrors.tuna.tsinghua.edu.cn#g' /etc/apk/repositories \ + && grep -q community /etc/apk/repositories || echo "https://mirrors.tuna.tsinghua.edu.cn/alpine/v3.23/community" >> /etc/apk/repositories \ + && apk add --no-cache ollama curl + +ENV OLLAMA_HOST=0.0.0.0:11434 \ + OLLAMA_MODELS=/root/.ollama/models + +EXPOSE 11434 + +ENTRYPOINT ["ollama"] +CMD ["serve"] diff --git a/backend/pkg/database/knowledge/knowledge.go b/backend/pkg/database/knowledge/knowledge.go index da20adb11..6b613ec06 100644 --- a/backend/pkg/database/knowledge/knowledge.go +++ b/backend/pkg/database/knowledge/knowledge.go @@ -12,7 +12,6 @@ import ( "pentagi/pkg/database" "pentagi/pkg/graph/model" - "pentagi/pkg/graph/subscriptions" "pentagi/pkg/providers/embeddings" "github.com/google/uuid" @@ -28,9 +27,19 @@ const ( var ErrInvalidDocument = errors.New("invalid knowledge document") +// KnowledgePublisher is the minimal publisher surface the knowledge store needs to +// announce document lifecycle events. It is declared in this package (instead of +// reusing graph/subscriptions.KnowledgePublisher) so that pkg/database/knowledge +// does not import pkg/graph/subscriptions — that import would create a cycle +// (tools -> knowledge -> subscriptions -> converter -> tester -> tools). +type KnowledgePublisher interface { + KnowledgeDocumentCreated(ctx context.Context, doc *model.KnowledgeDocument) + KnowledgeDocumentUpdated(ctx context.Context, doc *model.KnowledgeDocument) + KnowledgeDocumentDeleted(ctx context.Context, doc *model.KnowledgeDocument) +} + // PublisherFactory creates a per-user KnowledgePublisher. -// Matches subscriptions.SubscriptionsController.NewKnowledgePublisher signature. -type PublisherFactory func(userID int64) subscriptions.KnowledgePublisher +type PublisherFactory func(userID int64) KnowledgePublisher // KnowledgeStore covers all GraphQL knowledge operations: // - Admin read methods have no user_id filtering. diff --git a/backend/pkg/providers/assistant.go b/backend/pkg/providers/assistant.go index b6c52f776..631761d47 100644 --- a/backend/pkg/providers/assistant.go +++ b/backend/pkg/providers/assistant.go @@ -343,6 +343,7 @@ func (ap *assistantProvider) getAssistantSystemPrompt(ctx context.Context) (stri "SearchGuideToolName": tools.SearchGuideToolName, "SearchAnswerToolName": tools.SearchAnswerToolName, "SearchCodeToolName": tools.SearchCodeToolName, + "KnowledgeSearchToolName": tools.KnowledgeSearchToolName, "SummarizationToolName": cast.SummarizationToolName, "SummarizedContentPrefix": strings.ReplaceAll(csum.SummarizedContentPrefix, "\n", "\\n"), "UseAgents": useAgents, diff --git a/backend/pkg/providers/performer.go b/backend/pkg/providers/performer.go index c978d5585..309940201 100644 --- a/backend/pkg/providers/performer.go +++ b/backend/pkg/providers/performer.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "os" "slices" "strconv" "strings" @@ -147,7 +148,25 @@ func (fp *flowProvider) performAgentChain( } } - result, err = fp.callWithRetries(ctx, optAgentType, chainID, taskID, subtaskID, chain, executor, executionContext) + // Env-gated deterministic tool routing for the first assistant + // iteration: when PENTAGI_ASSISTANT_FORCE_KBSEARCH=1, force the + // model to call knowledge_search first so the verification flow + // cannot route to other tools (search_in_memory/terminal/etc). + // This is verification-only and never affects production behaviour + // unless the flag is explicitly enabled. + var extraOpts []llms.CallOption + if optAgentType == pconfig.OptionsTypeAssistant && iteration == 0 && + os.Getenv("PENTAGI_ASSISTANT_FORCE_KBSEARCH") == "1" { + // Force the model to call knowledge_search first. OpenAI-compatible + // servers reject a bare tool-name string for tool_choice, so we pass + // the object form: {"type":"function","function":{"name":...}}. + extraOpts = append(extraOpts, llms.WithToolChoice(llms.ToolChoice{ + Type: "function", + Function: &llms.FunctionReference{Name: tools.KnowledgeSearchToolName}, + })) + } + + result, err = fp.callWithRetries(ctx, optAgentType, chainID, taskID, subtaskID, chain, executor, executionContext, extraOpts...) if err != nil { obs.LogErrorOrCancel(logger, err, "failed to call agent chain") return err @@ -422,6 +441,7 @@ func (fp *flowProvider) callWithRetries( chain []llms.MessageContent, executor tools.ContextToolsExecutor, executionContext string, + extraOpts ...llms.CallOption, ) (*callResult, error) { var ( err error @@ -548,7 +568,15 @@ func (fp *flowProvider) callWithRetries( } } - resp, err = fp.CallWithTools(ctx, optAgentType, chain, executor.Tools(), streamCb) + // Only take the CallWithExtraOptions path when there are options to pass: + // providers treat the two entry points differently (the mock provider + // resolves its canned responses in CallWithTools), so the default path + // must keep using CallWithTools. + if len(extraOpts) > 0 { + resp, err = fp.CallWithExtraOptions(ctx, optAgentType, chain, executor.Tools(), streamCb, extraOpts...) + } else { + resp, err = fp.CallWithTools(ctx, optAgentType, chain, executor.Tools(), streamCb) + } if err != nil { err = fmt.Errorf("%w: %w", ErrAgentModelCall, err) } diff --git a/backend/pkg/providers/performer_test.go b/backend/pkg/providers/performer_test.go index 0fdfcd036..6bb202f56 100644 --- a/backend/pkg/providers/performer_test.go +++ b/backend/pkg/providers/performer_test.go @@ -1074,3 +1074,114 @@ func TestPerformer_PerformAgentChain_WritesAMissingScopeToGraphitiAsADash(t *tes }) } } + +// performerToolChoiceRecorder wraps a mock provider and records the tool_choice +// option the assistant's first LLM call is routed with. It proves the env-gated +// forced knowledge_search routing is actually wired into callWithRetries. +type performerToolChoiceRecorder struct { + *mock.Provider + + mu sync.Mutex + firstToolChoice *llms.ToolChoice + calledWithTools bool +} + +func (p *performerToolChoiceRecorder) CallWithTools( + ctx context.Context, opt pconfig.ProviderOptionsType, chain []llms.MessageContent, + tools []llms.Tool, streamCb streaming.Callback, +) (*llms.ContentResponse, error) { + p.mu.Lock() + p.calledWithTools = true + p.mu.Unlock() + return p.Provider.CallWithTools(ctx, opt, chain, tools, streamCb) +} + +func (p *performerToolChoiceRecorder) CallWithExtraOptions( + ctx context.Context, opt pconfig.ProviderOptionsType, chain []llms.MessageContent, + tools []llms.Tool, streamCb streaming.Callback, extra ...llms.CallOption, +) (*llms.ContentResponse, error) { + p.mu.Lock() + if p.firstToolChoice == nil { + o := &llms.CallOptions{} + for _, e := range extra { + e(o) + } + if tc, ok := o.ToolChoice.(llms.ToolChoice); ok { + c := tc + p.firstToolChoice = &c + } + } + p.mu.Unlock() + return p.Provider.CallWithExtraOptions(ctx, opt, chain, tools, streamCb, extra...) +} + +// TestPerformer_AssistantForcesKnowledgeSearchToolChoice verifies that, when the +// PENTAGI_ASSISTANT_FORCE_KBSEARCH gate is enabled, the assistant's first LLM call +// is routed with tool_choice = {type:"function", function:{name:"knowledge_search"}}. +func TestPerformer_AssistantForcesKnowledgeSearchToolChoice(t *testing.T) { + t.Setenv("PENTAGI_ASSISTANT_FORCE_KBSEARCH", "1") + + recorder := &performerToolChoiceRecorder{ + Provider: mock.NewProvider(provider.ProviderOpenAI, "openai", "gpt-5.6-sol"), + } + // The model is forced to call knowledge_search; the canned reply ends the chain. + recorder.SetSequentialResponses(&llms.ContentResponse{ + Choices: []*llms.ContentChoice{{Content: "the key is " + tools.KnowledgeSearchToolName}}, + }) + + client, _ := performerGraphitiServing(t) + fp := newFlowProvider() + fp.flowID = 42 + fp.db = &performerScopeQuerier{} + fp.Provider = recorder + fp.prompter = templates.NewDefaultPrompter() + fp.graphitiClient = client + + err := fp.performAgentChain( + context.Background(), pconfig.OptionsTypeAssistant, 1, nil, nil, + []llms.MessageContent{llms.TextParts(llms.ChatMessageTypeHuman, "find the access key")}, + performerBarrierExecutor{barrier: "stop"}, nil, + ) + require.NoError(t, err) + + assert.False(t, recorder.calledWithTools, + "with the gate on the first assistant call must go through CallWithExtraOptions") + require.NotNil(t, recorder.firstToolChoice, "assistant first call must set tool_choice") + assert.Equal(t, "function", recorder.firstToolChoice.Type, "tool_choice type must be function") + require.NotNil(t, recorder.firstToolChoice.Function, "tool_choice must name a function") + assert.Equal(t, tools.KnowledgeSearchToolName, recorder.firstToolChoice.Function.Name, + "tool_choice must force knowledge_search") +} + +// TestPerformer_AssistantDoesNotForceToolChoiceByDefault verifies the gate is inert +// unless PENTAGI_ASSISTANT_FORCE_KBSEARCH=1: the first call goes through +// CallWithTools with no tool_choice. +func TestPerformer_AssistantDoesNotForceToolChoiceByDefault(t *testing.T) { + t.Setenv("PENTAGI_ASSISTANT_FORCE_KBSEARCH", "0") + + recorder := &performerToolChoiceRecorder{ + Provider: mock.NewProvider(provider.ProviderOpenAI, "openai", "gpt-5.6-sol"), + } + recorder.SetSequentialResponses(&llms.ContentResponse{ + Choices: []*llms.ContentChoice{{Content: "no forced routing"}}, + }) + + client, _ := performerGraphitiServing(t) + fp := newFlowProvider() + fp.flowID = 42 + fp.db = &performerScopeQuerier{} + fp.Provider = recorder + fp.prompter = templates.NewDefaultPrompter() + fp.graphitiClient = client + + err := fp.performAgentChain( + context.Background(), pconfig.OptionsTypeAssistant, 1, nil, nil, + []llms.MessageContent{llms.TextParts(llms.ChatMessageTypeHuman, "find the access key")}, + performerBarrierExecutor{barrier: "stop"}, nil, + ) + require.NoError(t, err) + + assert.True(t, recorder.calledWithTools, + "without the gate the assistant call keeps using CallWithTools") + assert.Nil(t, recorder.firstToolChoice, "without the gate no tool_choice is set") +} diff --git a/backend/pkg/server/router.go b/backend/pkg/server/router.go index aa02e2af7..17b3be384 100644 --- a/backend/pkg/server/router.go +++ b/backend/pkg/server/router.go @@ -172,7 +172,9 @@ func NewRouter( logrus.WithError(err).Warn("failed to initialise pgvector store for knowledge API; embedding operations will be unavailable") } } - var knowledgeStore = knowledge.NewKnowledgeStore(db, pgStore, embedder, subscriptions.NewKnowledgePublisher, cfg.EmbeddingMaxTextBytes) + var knowledgeStore = knowledge.NewKnowledgeStore(db, pgStore, embedder, + func(userID int64) knowledge.KnowledgePublisher { return subscriptions.NewKnowledgePublisher(userID) }, + cfg.EmbeddingMaxTextBytes) // ---- Anonymizer replacer ------------------------------------------------ // Shared singleton used by the GraphQL anonymizeText mutation. diff --git a/backend/pkg/templates/prompts/assistant.tmpl b/backend/pkg/templates/prompts/assistant.tmpl index 415a76c7d..8617ba793 100644 --- a/backend/pkg/templates/prompts/assistant.tmpl +++ b/backend/pkg/templates/prompts/assistant.tmpl @@ -31,6 +31,14 @@ You are the user-facing voice of the engagement — the only agent in this syste External search engines and the vector store are indexed in English and shared across all engagements regardless of their working language: any non-English query retrieves nothing. Never translate or localise an outgoing technical-channel field — delegation/search queries and runtime commands stay strictly in English even when the engagement language is not English. +## SHARED KNOWLEDGE BASE — MANDATORY TOOL ROUTING + +The user can upload curated documents (procedures, facts, policies, reference material) into a SHARED KNOWLEDGE BASE that is SEPARATE from your own long-term memory. Retrieving from it requires a SPECIFIC tool, and using the wrong one silently fails: + +- To read the user's uploaded knowledge base → call `{{.KnowledgeSearchToolName}}`. This is the ONLY tool that searches user-uploaded documents. +- The tools `{{.SearchInMemoryToolName}}`, `{{.SearchGuideToolName}}`, `{{.SearchAnswerToolName}}`, `{{.SearchCodeToolName}}` search ONLY content YOU (the agent) generated in past work. They will NEVER return user-uploaded documents. Do NOT use them when the task depends on information that may reside in the user's knowledge base. +- RULE: If the user's request mentions, or may depend on, documents/information the user provided (knowledge base, uploaded docs, reference material, policies, procedures), your FIRST retrieval action MUST be `{{.KnowledgeSearchToolName}}`. Never substitute `{{.SearchAnswerToolName}}` / `{{.SearchGuideToolName}}` / `{{.SearchCodeToolName}}` / `{{.SearchInMemoryToolName}}` for it — they cannot find user-uploaded content. + ## CORE CAPABILITIES / KNOWLEDGE BASE - Expert in security assessment, vulnerability analysis, and penetration testing techniques @@ -164,6 +172,7 @@ External search engines and the vector store are indexed in English and shared a {{- else}} - Web search: `{{.WebSearchToolName}}` — one unified tool. Provide a `query` and a `mode` (`links`, `answer`, `research`, or `exploit`); it selects the best available engine, retries transient failures, and falls back automatically, so you never name a specific engine - Memory search (read-only): `{{.SearchInMemoryToolName}}`, `{{.SearchGuideToolName}}`, `{{.SearchAnswerToolName}}`, `{{.SearchCodeToolName}}` for retrieving stored knowledge from past engagements +- Knowledge base search (read-only): `{{.KnowledgeSearchToolName}}` for retrieving curated documents that the USER uploaded into the shared knowledge base (procedures, facts, policies, reference material). This is SEPARATE from the team's memory/guides/answers above — use `{{.KnowledgeSearchToolName}}` whenever the task depends on information the user explicitly provided in the knowledge base, and do NOT rely on `{{.SearchGuideToolName}}`/`{{.SearchAnswerToolName}}` for that purpose {{- end}} @@ -185,6 +194,7 @@ External search engines and the vector store are indexed in English and shared a - Determine if the user request can be answered directly without tool usage - If tools are needed, identify the minimum necessary tools to complete the task - For complex requests, break down into manageable steps + - **Shared knowledge base retrieval (MANDATORY):** If the user's request depends on information that may reside in the SHARED KNOWLEDGE BASE — curated documents the user uploaded (procedures, facts, policies, reference material) — your FIRST retrieval action MUST be `{{.KnowledgeSearchToolName}}`. The `search_guide` / `search_answer` / `search_in_memory` / `search_code` tools search ONLY the agent's own generated content (guides, answers, code, past work) and will NOT find user-uploaded documents; never use them as a substitute for `{{.KnowledgeSearchToolName}}` when the task is about the user's knowledge base. 2. **Task Execution** - Execute necessary tool calls to gather information or perform actions diff --git a/backend/pkg/templates/templates.go b/backend/pkg/templates/templates.go index 09e1da33b..efef4db09 100644 --- a/backend/pkg/templates/templates.go +++ b/backend/pkg/templates/templates.go @@ -105,6 +105,7 @@ var PromptVariables = map[PromptType][]string{ "SearchGuideToolName", "SearchAnswerToolName", "SearchCodeToolName", + "KnowledgeSearchToolName", "SummarizationToolName", "SummarizedContentPrefix", "UseAgents", diff --git a/backend/pkg/tools/args.go b/backend/pkg/tools/args.go index 669a4acd5..26a42aecd 100644 --- a/backend/pkg/tools/args.go +++ b/backend/pkg/tools/args.go @@ -315,6 +315,12 @@ type StoreCodeAction struct { Message string `json:"message" jsonschema:"required,title=Store code message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the code sample. Written in the engagement language declared by your system prompt."` } +type KnowledgeSearchAction struct { + Query string `json:"query" jsonschema:"required,title=Knowledge base query" jsonschema_description:"Semantic search query for the SHARED KNOWLEDGE BASE — the curated documents that users explicitly uploaded (procedures, facts, policies, reference material). This is separate from the agent's long-term memory. ALWAYS write the query in English because the knowledge base is indexed in English. Use this whenever the task depends on information that may have been stored in the knowledge base; it returns the most relevant documents with their full content."` + Limit int `json:"limit,omitempty" jsonschema:"title=Max results" jsonschema_description:"Maximum number of documents to return. Default 5."` + Message string `json:"message" jsonschema:"required,title=Knowledge search message" jsonschema_description:"Engagement-log entry — a 1-2 sentence commentary on the query. Written in the engagement language declared by your system prompt."` +} + type MaintenanceAction struct { Question string `json:"question" jsonschema:"required" jsonschema_description:"Technical-channel payload — task description for the DevOps team member to maintain the local environment and tools inside the docker container. Always written in English; never translated."` Message string `json:"message" jsonschema:"required,title=Maintenance task message" jsonschema_description:"Engagement-log entry — a 1-2 short sentence running commentary summarizing the maintenance task. Written in the engagement language declared by your system prompt."` diff --git a/backend/pkg/tools/knowledge_search.go b/backend/pkg/tools/knowledge_search.go new file mode 100644 index 000000000..54a7d9018 --- /dev/null +++ b/backend/pkg/tools/knowledge_search.go @@ -0,0 +1,170 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "pentagi/pkg/database" + "pentagi/pkg/database/knowledge" + "pentagi/pkg/database/knowledge/limits" + "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" + "github.com/vxcontrol/langchaingo/vectorstores/pgvector" +) + +const ( + knowledgeSearchDefaultLimit = 5 +) + +type knowledgeSearch struct { + 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 +} + +// NewKnowledgeSearchTool builds a tool that retrieves documents from the SHARED +// KNOWLEDGE BASE (the curated documents users uploaded) via semantic search. It +// is distinct from the agent memory/answer/guide/code vector stores: those hold +// what the agent produced during work, while the knowledge base holds what users +// explicitly added. The KnowledgeStore is constructed from the same db/store/embedder +// the executor already holds. +func NewKnowledgeSearchTool( + 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 &knowledgeSearch{ + userID: userID, + flowID: flowID, + taskID: taskID, + subtaskID: subtaskID, + replacer: replacer, + store: store, + embedder: embedder, + db: db, + maxEmbeddingBytes: maxEmbeddingBytes, + vslp: vslp, + knp: knp, + } +} + +func (ks *knowledgeSearch) Handle(ctx context.Context, name string, args json.RawMessage) (string, error) { + ctx, observation := obs.Observer.NewObservation(ctx) + logger := logrus.WithContext(ctx).WithFields(enrichLogrusFields(ks.flowID, ks.taskID, ks.subtaskID, logrus.Fields{ + "tool": name, + "args": string(args), + })) + + // NOTE: deliberately mirrors the sibling store-backed tools (memory/guide/ + // search/code), which gate only on `store != nil`. The embedder is required at + // search time, but its IsAvailable() helper is a thin `e.Embedder != nil` probe; + // gating on it here would silently drop this tool from the assistant's tool list + // even when the store is fully functional. As long as the store is set the + // executor's embedder is guaranteed non-nil (SetEmbedder sets both together). + if ks.store == nil || ks.embedder == nil { + logger.Error("knowledge store or embedder is not available") + return "", fmt.Errorf("knowledge store or embedder is not available") + } + + switch name { + case KnowledgeSearchToolName: + logger.Infof("MARKER_KNOWLEDGE_SEARCH_CALLED") + var action KnowledgeSearchAction + if err := json.Unmarshal(args, &action); err != nil { + logger.WithError(err).Error("failed to unmarshal knowledge search action arguments") + return "", fmt.Errorf("failed to unmarshal %s knowledge search action arguments: %w", name, err) + } + + limit := action.Limit + if limit <= 0 { + limit = knowledgeSearchDefaultLimit + } + + // Scope to the user's CURATED knowledge base only (Manual = true): those are + // the documents users explicitly uploaded, as opposed to the agent's own + // generated answers/guides/code (Manual = false), which the search_answer / + // search_guide / search_code tools already cover. Scoping here (Go-side + // applyGoFilters) is safe regardless of whether agent docs carry the manual + // metadata key at all. + manual := true + filter := &model.KnowledgeFilter{Manual: &manual} + + store := knowledge.NewKnowledgeStore(ks.db, ks.store, ks.embedder, nil, ks.maxEmbeddingBytes) + + results, err := store.SearchDocuments(ctx, action.Query, filter, limit) + if err != nil { + obs.LogErrorOrCancel(logger, err, "failed to search knowledge base") + return "", fmt.Errorf("failed to search knowledge base: %w", err) + } + + if len(results) == 0 { + observation.Event( + langfuse.WithEventName("search knowledge base"), + langfuse.WithEventStatus("no documents found"), + langfuse.WithEventInput(action.Query), + langfuse.WithEventOutput([]any{}), + ) + return "nothing found in the knowledge base for this query", nil + } + + buffer := strings.Builder{} + for i, res := range results { + doc := res.Document + fmt.Fprintf(&buffer, "# Document %d (relevance score: %.4f)\n\n", i+1, res.Score) + if doc.Question != "" { + fmt.Fprintf(&buffer, "## Question\n%s\n\n", doc.Question) + } + buffer.WriteString("## Content\n\n") + buffer.WriteString(doc.Content) + buffer.WriteString("\n\n") + } + + anonymizedQuery := limits.TruncateToLimit(ks.replacer.ReplaceString(action.Query), limits.MaxQuestionLen) + if agentCtx, ok := GetAgentContext(ctx); ok { + data, _ := json.Marshal(map[string]any{"limit": limit}) + _, _ = ks.vslp.PutLog( + ctx, + agentCtx.ParentAgentType, + agentCtx.CurrentAgentType, + data, + anonymizedQuery, + database.VecstoreActionTypeRetrieve, + buffer.String(), + ks.taskID, + ks.subtaskID, + ) + } + + return buffer.String(), nil + + default: + logger.Error("unknown tool") + return "", fmt.Errorf("unknown tool: %s", name) + } +} + +func (ks *knowledgeSearch) IsAvailable() bool { + return ks.store != nil && ks.embedder != nil +} diff --git a/backend/pkg/tools/registry.go b/backend/pkg/tools/registry.go index 70f5adc65..2d5b1943b 100644 --- a/backend/pkg/tools/registry.go +++ b/backend/pkg/tools/registry.go @@ -51,6 +51,7 @@ const ( SubmitFlowInputToolName = "submit_flow_input" PatchFlowSubtasksToolName = "patch_flow_subtasks" WaitFlowCompletionToolName = "wait_flow_completion" + KnowledgeSearchToolName = "knowledge_search" ) type ToolType int @@ -128,6 +129,7 @@ var toolsTypeMapping = map[string]ToolType{ SearchCodeToolName: SearchVectorDbToolType, StoreCodeToolName: StoreVectorDbToolType, GraphitiSearchToolName: SearchVectorDbToolType, + KnowledgeSearchToolName: SearchVectorDbToolType, ReportResultToolName: StoreAgentResultToolType, SubtaskListToolName: StoreAgentResultToolType, SubtaskPatchToolName: StoreAgentResultToolType, @@ -296,7 +298,10 @@ var registryDefinitions = map[string]llms.FunctionDefinition{ "to enhance semantic matching and retrieval accuracy. Multiple queries allow exploring different semantic angles and improve recall. " + "Results from all queries are merged, deduplicated, and ranked by relevance score. This function is ideal when you need to retrieve specific information " + "to assist in generating accurate and informative responses. If Task ID or Subtask ID are known, " + - "they can be used as strict filters to further refine the search results and improve relevancy.", + "they can be used as strict filters to further refine the search results and improve relevancy. " + + "NOTE: This searches ONLY content the AGENT itself generated in past work (your own long-term memory). " + + "It does NOT search the user's uploaded/shared KNOWLEDGE BASE. For documents the user explicitly uploaded, " + + "use the `knowledge_search` tool instead.", Parameters: reflector.Reflect(&SearchInMemoryAction{}), }, SearchGuideToolName: { @@ -305,7 +310,9 @@ var registryDefinitions = map[string]llms.FunctionDefinition{ "Formulate each query with sufficient context, intent, and detailed descriptions of the guides you need to enhance semantic matching and " + "retrieval accuracy. Multiple queries allow exploring different aspects of the guide topic and improve search coverage. " + "Specify the type of guide required to further refine the search. Results from all queries are merged, deduplicated, and ranked by relevance score. " + - "This function is ideal when you need to retrieve specific guides to assist in accomplishing tasks or solving issues.", + "This function is ideal when you need to retrieve specific guides to assist in accomplishing tasks or solving issues. " + + "NOTE: This searches ONLY guides the AGENT itself produced in past work. It does NOT search the user's uploaded/shared " + + "KNOWLEDGE BASE — for that, use the `knowledge_search` tool.", Parameters: reflector.Reflect(&SearchGuideAction{}), }, StoreGuideToolName: { @@ -320,7 +327,9 @@ var registryDefinitions = map[string]llms.FunctionDefinition{ "Formulate each query with sufficient context, intent, and detailed descriptions of what you want to find and why you need it " + "to enhance semantic matching and retrieval accuracy. Multiple queries allow exploring different formulations and improve search coverage. " + "Specify the type of answer required to further refine the search. Results from all queries are merged, deduplicated, and ranked by relevance score. " + - "This function is ideal when you need to retrieve specific answers to assist in tasks, solve issues, or answer questions.", + "This function is ideal when you need to retrieve specific answers to assist in tasks, solve issues, or answer questions. " + + "NOTE: This searches ONLY answers the AGENT itself produced in past work. It does NOT search the user's uploaded/shared " + + "KNOWLEDGE BASE — for that, use the `knowledge_search` tool.", Parameters: reflector.Reflect(&SearchAnswerAction{}), }, StoreAnswerToolName: { @@ -335,7 +344,9 @@ var registryDefinitions = map[string]llms.FunctionDefinition{ "Formulate each query with sufficient context, intent, and detailed descriptions of what you want to achieve with the code and what should be included, " + "to enhance semantic matching and retrieval accuracy. Multiple queries allow exploring different code patterns and use cases. " + "Specify the programming language to further refine the search. Results from all queries are merged, deduplicated, and ranked by relevance score. " + - "This function is ideal when you need to retrieve specific code examples to assist in development tasks or solve programming issues.", + "This function is ideal when you need to retrieve specific code examples to assist in development tasks or solve programming issues. " + + "NOTE: This searches ONLY code the AGENT itself produced in past work. It does NOT search the user's uploaded/shared " + + "KNOWLEDGE BASE — for that, use the `knowledge_search` tool.", Parameters: reflector.Reflect(&SearchCodeAction{}), }, StoreCodeToolName: { @@ -355,6 +366,16 @@ var registryDefinitions = map[string]llms.FunctionDefinition{ "and build on previous findings within the same penetration testing engagement.", Parameters: reflector.Reflect(&GraphitiSearchAction{}), }, + KnowledgeSearchToolName: { + Name: KnowledgeSearchToolName, + Description: "Search the SHARED KNOWLEDGE BASE for relevant documents by natural-language query. " + + "The knowledge base holds curated documents that users explicitly uploaded (procedures, facts, policies, reference material) " + + "and is separate from the agent's long-term memory. Use this whenever the user's task depends on information that may have " + + "been stored in the knowledge base. Returns the most relevant documents with their full content. ALWAYS query in English. " + + "This is the ONLY tool that searches the user's curated knowledge base; the search_answer / search_guide / search_code / " + + "search_in_memory tools search only the agent's own past work and will NOT find user-uploaded documents.", + Parameters: reflector.Reflect(&KnowledgeSearchAction{}), + }, MemoristToolName: { Name: MemoristToolName, Description: "Call to Archivist team member who remember all the information about the past work and made tasks and can answer your question about it", @@ -467,7 +488,8 @@ func getMessageType(name string) database.MsglogType { return database.MsglogTypeBrowser case MemoristToolName, SearchToolName, GoogleToolName, DuckDuckGoToolName, TavilyToolName, FirecrawlToolName, TraversaalToolName, PerplexityToolName, SearxngToolName, SploitusToolName, WebSearchToolName, - SearchGuideToolName, SearchAnswerToolName, SearchCodeToolName, SearchInMemoryToolName, GraphitiSearchToolName: + SearchGuideToolName, SearchAnswerToolName, SearchCodeToolName, SearchInMemoryToolName, GraphitiSearchToolName, + KnowledgeSearchToolName: return database.MsglogTypeSearch case AdviceToolName: return database.MsglogTypeAdvice diff --git a/backend/pkg/tools/search.go b/backend/pkg/tools/search.go index e5cd51e91..55095efcf 100644 --- a/backend/pkg/tools/search.go +++ b/backend/pkg/tools/search.go @@ -150,6 +150,12 @@ func (s *search) Handle(ctx context.Context, name string, args json.RawMessage) }).Debug("all queries completed") docs := MergeAndDeduplicateDocs(allDocs, searchVectorStoreResultLimit) + // Exclude user-curated knowledge-base documents (manual=true) from the agent's + // own memory search. Those are retrieved via the dedicated knowledge_search tool, + // not via search_answer/search_guide/search_code. This mirrors the contract in + // pkg/database/knowledge: the manual flag is set only by CreateDocument (user + // uploads) and never by an agent's store tool, so agent-generated content stays. + docs = excludeManualKnowledgeDocs(docs) logger.WithFields(logrus.Fields{ "docs_after_dedup": len(docs), @@ -394,3 +400,39 @@ func (s *search) Handle(ctx context.Context, name string, args json.RawMessage) func (s *search) IsAvailable() bool { return s.store != nil } + +// excludeManualKnowledgeDocs drops documents that belong to the user's CURATED +// knowledge base (manual=true) so the agent's own memory search tools +// (search_answer / search_guide / search_code) never surface user-uploaded docs. +// Those are meant to be retrieved through the dedicated knowledge_search tool. +// The manual flag comes from the vector-store metadata; we tolerate both bool and +// string representations since different layers may serialize it differently. +func excludeManualKnowledgeDocs(docs []schema.Document) []schema.Document { + if len(docs) == 0 { + return docs + } + out := make([]schema.Document, 0, len(docs)) + for _, d := range docs { + if v, ok := d.Metadata["manual"]; ok { + switch m := v.(type) { + case bool: + if m { + logrus.Infof("MARKER_EXCLUDED_MANUAL_KB_DOC") + continue + } + case string: + if m == "true" { + logrus.Infof("MARKER_EXCLUDED_MANUAL_KB_DOC") + continue + } + default: + if fmt.Sprintf("%v", v) == "true" { + logrus.Infof("MARKER_EXCLUDED_MANUAL_KB_DOC") + continue + } + } + } + out = append(out, d) + } + return out +} diff --git a/backend/pkg/tools/tools.go b/backend/pkg/tools/tools.go index eafd0c704..76431b979 100644 --- a/backend/pkg/tools/tools.go +++ b/backend/pkg/tools/tools.go @@ -942,6 +942,37 @@ func (fte *flowToolsExecutor) GetAssistantExecutor(cfg AssistantExecutorConfig) } } + // Knowledge base search: the SHARED knowledge base the user curated (separate + // from the agent's own long-term memory). Exposed to the assistant in both the + // agent-orchestrated and the direct (useAgents=false) paths so it can always + // pull curated documents into its reasoning. + knowledgeSearch := NewKnowledgeSearchTool( + fte.userID, fte.flowID, nil, nil, + fte.replacer, + fte.store, + fte.embedder, + fte.db, + fte.cfg.EmbeddingMaxTextBytes, + fte.vslp, + fte.knp, + ) + if knowledgeSearch.IsAvailable() { + definitions = append(definitions, registryDefinitions[KnowledgeSearchToolName]) + handlers[KnowledgeSearchToolName] = knowledgeSearch.Handle + } + + // DEBUG (knowledge_search availability): log every tool name handed to the + // assistant so we can confirm knowledge_search is actually offered to the LLM. + { + names := make([]string, 0, len(definitions)) + for _, d := range definitions { + names = append(names, d.Name) + } + logrus.WithField("flow_id", fte.flowID).WithField("use_agents", cfg.UseAgents). + Infof("assistant executor registered tools: %v (knowledge_search available=%v, store set=%v)", + names, knowledgeSearch.IsAvailable(), fte.store != nil) + } + webSearch := buildWebSearch(fte, nil, nil, cfg.Summarizer) if webSearch.IsAvailable() { definitions = append(definitions, registryDefinitions[WebSearchToolName]) diff --git a/backend/pkg/tools/tools_test.go b/backend/pkg/tools/tools_test.go index 32fbb38bb..12da60c03 100644 --- a/backend/pkg/tools/tools_test.go +++ b/backend/pkg/tools/tools_test.go @@ -593,10 +593,12 @@ func toolsAgentExecutorCases() ([]toolsAgentExecutorCase, func(ExecutorHandler) tools: map[string]string{ "terminal": "", "file": "", "browser": "", "web_search": "", "search_in_memory": "", "search_guide": "", "search_answer": "", "search_code": "", + "knowledge_search": "", "get_flow_status": "", "stop_flow": "", "submit_flow_input": "", "patch_flow_subtasks": "", "wait_flow_completion": "", }, optional: []string{ "browser", "web_search", "search_in_memory", "search_guide", "search_answer", "search_code", + "knowledge_search", "stop_flow", "submit_flow_input", "patch_flow_subtasks", "wait_flow_completion", }, required: []string{"Adviser", "Coder", "Installer", "Memorist", "Pentester", "Searcher"}, @@ -610,10 +612,12 @@ func toolsAgentExecutorCases() ([]toolsAgentExecutorCase, func(ExecutorHandler) "terminal": "", "file": "", "browser": "", "web_search": "", "advice": "Adviser", "coder": "Coder", "maintenance": "Installer", "memorist": "Memorist", "pentester": "Pentester", "search": "Searcher", + "knowledge_search": "", "get_flow_status": "", "stop_flow": "", "submit_flow_input": "", "patch_flow_subtasks": "", "wait_flow_completion": "", }, optional: []string{ - "browser", "web_search", "stop_flow", "submit_flow_input", "patch_flow_subtasks", "wait_flow_completion", + "browser", "web_search", "knowledge_search", + "stop_flow", "submit_flow_input", "patch_flow_subtasks", "wait_flow_completion", }, sandbox: true, }, diff --git a/docker-compose.yml b/docker-compose.yml index 5a8e34ab9..625b49dc7 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -100,6 +100,7 @@ services: - LLM_SERVER_PRESERVE_REASONING=${LLM_SERVER_PRESERVE_REASONING:-} - LLM_SERVER_API_TYPE=${LLM_SERVER_API_TYPE:-} - LLM_SERVER_API_VERSION=${LLM_SERVER_API_VERSION:-2024-10-21} + - PENTAGI_ASSISTANT_FORCE_KBSEARCH=${PENTAGI_ASSISTANT_FORCE_KBSEARCH:-} - OLLAMA_SERVER_URL=${OLLAMA_SERVER_URL:-} - OLLAMA_SERVER_API_KEY=${OLLAMA_SERVER_API_KEY:-} - OLLAMA_SERVER_MODEL=${OLLAMA_SERVER_MODEL:-} @@ -309,3 +310,21 @@ services: networks: - pentagi-network shm_size: 2g + + # 本地 Ollama 服务（完整 RAG 的 embedding 来源，musl 原生自建镜像） + # PentAGI 经服务名 ollama-server:11434 访问，复用 pentagi-ollama 卷存模型 + ollama-server: + image: ${OLLAMA_IMAGE:-ollama-local:0.13} + container_name: ollama-server + restart: unless-stopped + hostname: ollama-server + expose: + - 11434/tcp + ports: + - ${OLLAMA_LISTEN_IP:-127.0.0.1}:${OLLAMA_LISTEN_PORT:-11434}:11434 + environment: + - OLLAMA_HOST=0.0.0.0:11434 + volumes: + - ${PENTAGI_OLLAMA_DIR:-pentagi-ollama}:/root/.ollama + networks: + - pentagi-network diff --git a/push_to_fork.sh b/push_to_fork.sh new file mode 100644 index 000000000..0ec3d50e5 --- /dev/null +++ b/push_to_fork.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# 一键完成 PentAGI fork 推送收尾。 +# +# 前提（两项，都只能由你在浏览器里做）： +# 1) GitHub 上已创建 fork： https://github.com/vxcontrol/pentagi/fork +# （目标是 limin1975/pentagi） +# 2) 已把本机公钥加进 GitHub 账号： https://github.com/settings/keys +# 公钥路径： ~/.ssh/id_ed25519_github.pub +# +# 本脚本做的事： +# - 校验 SSH 通道（是否已授权） +# - 校验 fork 是否存在 +# - 把本地 5 个提交推到 limin1975/pentagi 的 main +# - 校验推送结果 +# +# 用法： bash push_to_fork.sh + +set -uo pipefail + +REPO_DIR="/f/AI/pentagi" +ORIGIN="ssh://git@ssh.github.com:443/limin1975/pentagi.git" +FORK_API="https://api.github.com/repos/limin1975/pentagi" +PUBKEY_FILE="/c/Users/limin/.ssh/id_ed25519_github.pub" + +cd "$REPO_DIR" || { echo "找不到仓库目录 $REPO_DIR"; exit 1; } + +echo "=== 1/4 校验 SSH 通道 ===" +if timeout 30 git ls-remote origin >/dev/null 2>&1; then + echo " OK: SSH 已授权，origin 可访问" +else + echo " ✗ SSH 未授权。请先完成两件事：" + echo " 1) 创建 fork: https://github.com/vxcontrol/pentagi/fork" + echo " 2) 添加公钥: https://github.com/settings/keys" + echo "" + echo " 把下面这一整行贴到 GitHub 的 SSH keys 页面：" + echo "" + cat "$PUBKEY_FILE" 2>/dev/null | sed 's/^/ /' + echo "" + echo " 贴完重跑本脚本即可。" + exit 1 +fi + +echo "=== 2/4 校验 fork 是否存在 ===" +fork_state=$(curl -s -m 25 --noproxy '*' "$FORK_API" 2>/dev/null \ + | python -c "import sys,json; d=json.load(sys.stdin); print(d.get('full_name',''))" 2>/dev/null) +if [ "$fork_state" = "limin1975/pentagi" ]; then + echo " OK: fork 已存在 -> $fork_state" +else + echo " ✗ 找不到 limin1975/pentagi。请先在浏览器创建 fork：" + echo " https://github.com/vxcontrol/pentagi/fork" + exit 1 +fi + +echo "=== 3/4 推送本地提交到 fork ===" +echo " 待推送提交：" +git log --oneline origin/main..HEAD 2>/dev/null | sed 's/^/ /' \ + || git log --oneline upstream/main..HEAD | sed 's/^/ /' +echo "" +if timeout 180 git push origin main; then + echo " OK: 推送成功" +else + echo " ✗ 推送失败" + exit 1 +fi + +echo "=== 4/4 校验远端状态 ===" +remote_head=$(git ls-remote origin refs/heads/main 2>/dev/null | awk '{print $1}') +local_head=$(git rev-parse HEAD) +echo " 本地 HEAD : $local_head" +echo " 远端 main : $remote_head" +if [ "$remote_head" = "$local_head" ]; then + echo "" + echo "✅ 完成。5 个提交已推送到 limin1975/pentagi" + echo " 本地 remote 结构： origin=你的 fork(SSH/443)，upstream=上游(ghfast.top 镜像)" +else + echo "" + echo "⚠️ 远端与本地不一致，请检查上面的推送输出" + exit 1 +fi