From 45afdbdfbb6215c3a0d1935daeb554b78277be5c Mon Sep 17 00:00:00 2001
From: WangCham <651122857@qq.com>
Date: Thu, 17 Jul 2025 15:05:11 +0800
Subject: [PATCH] feat: rag pipeline backend

---
 pkg/entity/persistence/pipeline.py     |  3 ++-
 pkg/provider/runners/localagent.py     | 37 +++++++++++++++++++++++++-
 templates/default-pipeline-config.json |  3 ++-
 templates/metadata/pipeline/ai.yaml    |  8 ++++++
 4 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/pkg/entity/persistence/pipeline.py b/pkg/entity/persistence/pipeline.py
index 56e2cae9..a07470f3 100644
--- a/pkg/entity/persistence/pipeline.py
+++ b/pkg/entity/persistence/pipeline.py
@@ -20,7 +20,7 @@ class LegacyPipeline(Base):
     )
     for_version = sqlalchemy.Column(sqlalchemy.String(255), nullable=False)
     is_default = sqlalchemy.Column(sqlalchemy.Boolean, nullable=False, default=False)
-
+    knowledge_base_uuid = sqlalchemy.Column(sqlalchemy.String(255), nullable=True)
     stages = sqlalchemy.Column(sqlalchemy.JSON, nullable=False)
     config = sqlalchemy.Column(sqlalchemy.JSON, nullable=False)
 
@@ -43,3 +43,4 @@ class PipelineRunRecord(Base):
     started_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
     finished_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
     result = sqlalchemy.Column(sqlalchemy.JSON, nullable=False)
+    knowledge_base_uuid = sqlalchemy.Column(sqlalchemy.String(255), nullable=True)
diff --git a/pkg/provider/runners/localagent.py b/pkg/provider/runners/localagent.py
index 7d5e04c5..16d61e1b 100644
--- a/pkg/provider/runners/localagent.py
+++ b/pkg/provider/runners/localagent.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import json
 import typing
-
+from ...platform.types import message as platform_entities
 from .. import runner
 from ...core import entities as core_entities
 from .. import entities as llm_entities
@@ -15,9 +15,44 @@ class LocalAgentRunner(runner.RequestRunner):
     async def run(self, query: core_entities.Query) -> typing.AsyncGenerator[llm_entities.Message, None]:
         """运行请求"""
         pending_tool_calls = []
+            
 
         req_messages = query.prompt.messages.copy() + query.messages.copy() + [query.user_message]
 
+        
+        pipeline_uuid = query.pipeline_uuid
+        pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(pipeline_uuid)
+
+        try:
+            if pipeline and pipeline.pipeline_entity.knowledge_base_uuid is not None:
+                kb_id = pipeline.pipeline_entity.knowledge_base_uuid
+                kb= await self.ap.rag_mgr.load_knowledge_base(kb_id)
+        except Exception as e:
+            self.ap.logger.error(f'Failed to load knowledge base {kb_id}: {e}')
+            kb_id = None
+
+        if kb:
+            message = ''
+            for msg in query.message_chain:
+                if isinstance(msg, platform_entities.Plain):
+                    message += msg.text
+            result = await kb.retrieve(message)
+
+            if result:
+                rag_context = "\n\n".join(
+                    f"[{i+1}] {entry.metadata.get('text', '')}" for i, entry in enumerate(result)
+                )
+                rag_message = llm_entities.Message(
+                    role="user",
+                    content="The following are relevant context entries retrieved from the knowledge base. "
+                            "Please use them to answer the user's question. "
+                            "Respond in the same language as the user's input.\n\n" + rag_context
+                )
+                req_messages += [rag_message]
+
+
+
+
         # 首次请求
         msg = await query.use_llm_model.requester.invoke_llm(
             query,
diff --git a/templates/default-pipeline-config.json b/templates/default-pipeline-config.json
index 796c6356..d06e4661 100644
--- a/templates/default-pipeline-config.json
+++ b/templates/default-pipeline-config.json
@@ -44,7 +44,8 @@
                     "role": "system",
                     "content": "You are a helpful assistant."
                 }
-            ]
+            ],
+            "knowledge-base": ""
         },
         "dify-service-api": {
             "base-url": "https://api.dify.ai/v1",
diff --git a/templates/metadata/pipeline/ai.yaml b/templates/metadata/pipeline/ai.yaml
index 90732dc8..5f5dea5b 100644
--- a/templates/metadata/pipeline/ai.yaml
+++ b/templates/metadata/pipeline/ai.yaml
@@ -68,6 +68,13 @@ stages:
           zh_Hans: 除非您了解消息结构，否则请只使用 system 单提示词
         type: prompt-editor
         required: true
+      - name: knowledge-base
+        label:
+          en_US: Knowledge Base
+          zh_Hans: 知识库
+        type: knowledge-base-selector
+        required: false
+        default: ''
   - name: dify-service-api
     label:
       en_US: Dify Service API
@@ -298,3 +305,4 @@ stages:
         type: string
         required: false
         default: 'response'
+