mirror of
https://github.com/csunny/DB-GPT.git
synced 2026-07-16 17:15:22 +00:00
feat:knowledge base support agentic search
This commit is contained in:
aries_ckt
@@ -1,114 +0,0 @@
|
||||
"""Search Tools API Endpoints.
|
||||
|
||||
Provides REST API for knowledge base search tools, enabling the frontend
|
||||
search test panel and direct API access to kb_ls, kb_glob, kb_grep,
|
||||
kb_cat, and kb_semantic_search.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from dbgpt.component import SystemApp
|
||||
from dbgpt_serve.core import Result
|
||||
|
||||
from ..api.schemas import KbSearchRequest
|
||||
from ..config import SERVE_SERVICE_COMPONENT_NAME, ServeConfig
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
global_system_app: Optional[SystemApp] = None
|
||||
|
||||
|
||||
@router.post("/knowledge/{space_id}/tools/ls")
|
||||
async def kb_ls_endpoint(
|
||||
space_id: str,
|
||||
request: KbSearchRequest,
|
||||
) -> Result:
|
||||
"""List files and directories in a knowledge space."""
|
||||
from ..tools.kb_file_tools import kb_ls
|
||||
|
||||
result = await kb_ls(
|
||||
knowledge_id=space_id,
|
||||
path=request.path,
|
||||
offset=request.offset,
|
||||
limit=request.limit,
|
||||
)
|
||||
return Result.succ(result)
|
||||
|
||||
|
||||
@router.post("/knowledge/{space_id}/tools/glob")
|
||||
async def kb_glob_endpoint(
|
||||
space_id: str,
|
||||
request: KbSearchRequest,
|
||||
) -> Result:
|
||||
"""Search files by name pattern in a knowledge space."""
|
||||
from ..tools.kb_file_tools import kb_glob
|
||||
|
||||
result = await kb_glob(
|
||||
knowledge_id=space_id,
|
||||
pattern=request.query,
|
||||
limit=request.limit,
|
||||
offset=request.offset,
|
||||
)
|
||||
return Result.succ(result)
|
||||
|
||||
|
||||
@router.post("/knowledge/{space_id}/tools/grep")
|
||||
async def kb_grep_endpoint(
|
||||
space_id: str,
|
||||
request: KbSearchRequest,
|
||||
) -> Result:
|
||||
"""Search file contents by keyword in a knowledge space."""
|
||||
from ..tools.kb_file_tools import kb_grep
|
||||
|
||||
result = await kb_grep(
|
||||
knowledge_id=space_id,
|
||||
query=request.query,
|
||||
path=request.path,
|
||||
file_pattern=request.file_pattern,
|
||||
limit=request.limit,
|
||||
offset=request.offset,
|
||||
)
|
||||
return Result.succ(result)
|
||||
|
||||
|
||||
@router.post("/knowledge/{space_id}/tools/cat")
|
||||
async def kb_cat_endpoint(
|
||||
space_id: str,
|
||||
request: KbSearchRequest,
|
||||
) -> Result:
|
||||
"""Read file content from a knowledge space."""
|
||||
from ..tools.kb_file_tools import kb_cat
|
||||
|
||||
result = await kb_cat(
|
||||
knowledge_id=space_id,
|
||||
path=request.path,
|
||||
start_line=request.start_line,
|
||||
end_line=request.end_line,
|
||||
)
|
||||
return Result.succ(result)
|
||||
|
||||
|
||||
@router.post("/knowledge/{space_id}/tools/semantic_search")
|
||||
async def kb_semantic_search_endpoint(
|
||||
space_id: str,
|
||||
request: KbSearchRequest,
|
||||
) -> Result:
|
||||
"""Perform semantic search in a knowledge space."""
|
||||
from ..tools.semantic_search_tool import kb_semantic_search
|
||||
|
||||
result = await kb_semantic_search(
|
||||
knowledge_id=space_id,
|
||||
query=request.query,
|
||||
top_k=request.top_k,
|
||||
score_threshold=request.score_threshold,
|
||||
)
|
||||
return Result.succ(result)
|
||||
|
||||
|
||||
def init_search_endpoints(system_app: SystemApp, config: ServeConfig) -> None:
|
||||
"""Initialize the search tools endpoints."""
|
||||
global global_system_app
|
||||
global_system_app = system_app
|
||||
@@ -1,76 +0,0 @@
|
||||
"""Base class for Domain Knowledge Index.
|
||||
|
||||
Defines the ETL (Extract-Transform-Load) pipeline interface that each data source
|
||||
must implement. Each domain type (normal, git_repo, yuque, notion, etc.) provides
|
||||
its own indexing strategy through this abstraction.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional
|
||||
|
||||
from dbgpt.core import Chunk
|
||||
from dbgpt.rag.knowledge.base import Knowledge
|
||||
from dbgpt.storage.full_text.base import FullTextStoreBase
|
||||
from dbgpt.storage.knowledge_graph.base import KnowledgeGraphBase
|
||||
from dbgpt.storage.vector_store.base import VectorStoreBase
|
||||
|
||||
|
||||
class DomainKnowledgeIndex(ABC):
|
||||
"""Abstract base class for domain-specific knowledge indexing.
|
||||
|
||||
Each data source type (local documents, git repositories, yuque, notion, etc.)
|
||||
implements its own ETL pipeline by subclassing this and overriding the
|
||||
extract/transform/load methods.
|
||||
|
||||
The factory pattern (DomainKnowledgeIndexFactory) is used to instantiate the
|
||||
correct index based on the knowledge space's domain_type.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def extract(
|
||||
self,
|
||||
knowledge: Knowledge,
|
||||
chunk_parameter,
|
||||
**kwargs,
|
||||
) -> list[Chunk]:
|
||||
"""Extract knowledge chunks from the data source."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def transform(
|
||||
self,
|
||||
chunks: list[Chunk],
|
||||
**kwargs,
|
||||
) -> list[Chunk]:
|
||||
"""Transform knowledge chunks (enrichment, summarization, etc.)."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def load(
|
||||
self,
|
||||
chunks: list[Chunk],
|
||||
vector_store: Optional[VectorStoreBase] = None,
|
||||
full_text_store: Optional[FullTextStoreBase] = None,
|
||||
kg_store: Optional[KnowledgeGraphBase] = None,
|
||||
keywords: bool = True,
|
||||
max_chunks_once_load: int = 10,
|
||||
max_threads: int = 1,
|
||||
**kwargs,
|
||||
) -> list[Chunk]:
|
||||
"""Load knowledge chunks into storage backends."""
|
||||
raise NotImplementedError
|
||||
|
||||
async def clean(
|
||||
self,
|
||||
chunks: list[Chunk],
|
||||
node_ids: Optional[list[str]],
|
||||
with_keywords: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
"""Clean up indexed chunks from storage backends."""
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def domain_type(cls) -> str:
|
||||
"""Return the domain type identifier for this index."""
|
||||
raise NotImplementedError
|
||||
@@ -1,175 +0,0 @@
|
||||
"""Git Repository Index - ETL pipeline for Git repo knowledge bases."""
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
from collections import defaultdict
|
||||
from typing import List, Optional
|
||||
|
||||
from dbgpt.core import Chunk
|
||||
from dbgpt.rag.knowledge.base import Knowledge
|
||||
from dbgpt.storage.full_text.base import FullTextStoreBase
|
||||
from dbgpt.storage.knowledge_graph.base import KnowledgeGraphBase
|
||||
from dbgpt.storage.vector_store.base import VectorStoreBase
|
||||
from dbgpt_ext.rag import ChunkParameters
|
||||
from dbgpt_ext.rag.chunk_manager import ChunkManager
|
||||
|
||||
from .index import DomainGeneralIndex
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GitRepoIndex(DomainGeneralIndex):
|
||||
"""Git repository indexing pipeline with document-level summaries."""
|
||||
|
||||
async def extract(
|
||||
self,
|
||||
knowledge: Knowledge,
|
||||
chunk_parameter: ChunkParameters,
|
||||
extract_image: bool = False,
|
||||
**kwargs,
|
||||
) -> List[Chunk]:
|
||||
if not knowledge:
|
||||
raise ValueError("knowledge must be provided.")
|
||||
documents = await knowledge.aload()
|
||||
|
||||
# Check if GitRepoKnowledge has custom extract
|
||||
if hasattr(knowledge, "extract") and callable(
|
||||
getattr(knowledge, "extract", None)
|
||||
):
|
||||
try:
|
||||
result = knowledge.extract(documents, chunk_parameter)
|
||||
if result is not None:
|
||||
all_chunks = []
|
||||
for doc in result:
|
||||
if hasattr(doc, "chunks") and doc.chunks:
|
||||
for chunk in doc.chunks:
|
||||
chunk.metadata["chunk_id"] = chunk.chunk_id
|
||||
all_chunks.append(chunk)
|
||||
if all_chunks:
|
||||
return all_chunks
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"GitRepoKnowledge.extract() failed, falling back to ChunkManager: {e}"
|
||||
)
|
||||
|
||||
chunk_manager = ChunkManager(
|
||||
knowledge=knowledge, chunk_parameter=chunk_parameter
|
||||
)
|
||||
chunks = chunk_manager.split(documents)
|
||||
for chunk in chunks:
|
||||
chunk.metadata["chunk_id"] = chunk.chunk_id
|
||||
return chunks
|
||||
|
||||
async def transform(
|
||||
self,
|
||||
chunks: List[Chunk],
|
||||
image_extractor=None,
|
||||
summary_extractor=None,
|
||||
batch_size: int = 1,
|
||||
**kwargs,
|
||||
) -> List[Chunk]:
|
||||
transform_chunks = await super().transform(
|
||||
chunks,
|
||||
image_extractor=image_extractor,
|
||||
summary_extractor=None,
|
||||
batch_size=batch_size,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Add context prefix to each chunk
|
||||
for chunk in transform_chunks:
|
||||
doc_name = (chunk.metadata or {}).get("doc_name", "")
|
||||
file_path = (chunk.metadata or {}).get("file_path", "")
|
||||
context_prefix = ""
|
||||
if doc_name:
|
||||
context_prefix += f"[文件: {doc_name}]"
|
||||
if file_path:
|
||||
context_prefix += f" [路径: {file_path}]"
|
||||
if context_prefix:
|
||||
context_prefix = context_prefix.strip() + "\n"
|
||||
chunk.content = context_prefix + chunk.content
|
||||
|
||||
# Generate document-level summaries
|
||||
if summary_extractor:
|
||||
summary_chunks = await self._generate_doc_summaries(
|
||||
transform_chunks, summary_extractor
|
||||
)
|
||||
transform_chunks.extend(summary_chunks)
|
||||
|
||||
return transform_chunks
|
||||
|
||||
async def _generate_doc_summaries(
|
||||
self, chunks: List[Chunk], summary_extractor
|
||||
) -> List[Chunk]:
|
||||
doc_chunks = defaultdict(list)
|
||||
for chunk in chunks:
|
||||
if chunk.chunk_type == "image":
|
||||
continue
|
||||
doc_id = chunk.metadata.get("doc_id", "") if chunk.metadata else ""
|
||||
if doc_id:
|
||||
doc_chunks[doc_id].append(chunk)
|
||||
|
||||
summary_chunks = []
|
||||
for doc_id, doc_chunk_list in doc_chunks.items():
|
||||
try:
|
||||
full_text = "\n\n".join(
|
||||
c.content for c in doc_chunk_list if c.content
|
||||
)
|
||||
if not full_text.strip():
|
||||
continue
|
||||
max_chars = 30000
|
||||
if len(full_text) > max_chars:
|
||||
full_text = full_text[:max_chars] + "\n...(truncated)"
|
||||
|
||||
first_chunk = doc_chunk_list[0]
|
||||
file_type = (first_chunk.metadata or {}).get("file_type", "markdown")
|
||||
file_path = (first_chunk.metadata or {}).get("file_path", "")
|
||||
|
||||
if hasattr(summary_extractor, "generate_summary"):
|
||||
summary_text = await summary_extractor.generate_summary(
|
||||
content=full_text, file_type=file_type, file_path=file_path
|
||||
)
|
||||
else:
|
||||
summary_text = await summary_extractor.extract(text=full_text)
|
||||
if isinstance(summary_text, list):
|
||||
summary_text = summary_text[0] if summary_text else ""
|
||||
|
||||
if not summary_text:
|
||||
continue
|
||||
|
||||
doc_name = (first_chunk.metadata or {}).get("doc_name", "")
|
||||
summary_prefix = ""
|
||||
if doc_name:
|
||||
summary_prefix += f"[文件: {doc_name}]"
|
||||
if file_path:
|
||||
summary_prefix += f" [路径: {file_path}]"
|
||||
if summary_prefix:
|
||||
summary_text = summary_prefix.strip() + "\n" + summary_text
|
||||
|
||||
summary_metadata = {
|
||||
**(first_chunk.metadata or {}),
|
||||
"chunk_type": "summary",
|
||||
"doc_id": doc_id,
|
||||
}
|
||||
summary_chunks.append(
|
||||
Chunk(
|
||||
chunk_id=str(uuid.uuid4()),
|
||||
content=summary_text,
|
||||
metadata=summary_metadata,
|
||||
chunk_type="summary",
|
||||
summary=summary_text,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to generate summary for doc_id={doc_id}: {e}")
|
||||
continue
|
||||
|
||||
logger.info(
|
||||
f"Generated {len(summary_chunks)} document summaries "
|
||||
f"for {len(doc_chunks)} documents"
|
||||
)
|
||||
return summary_chunks
|
||||
|
||||
@classmethod
|
||||
def domain_type(cls) -> str:
|
||||
return "git_repo"
|
||||
@@ -1,137 +0,0 @@
|
||||
"""General Domain Knowledge Index - ETL pipeline for local documents."""
|
||||
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from dbgpt.core import Chunk
|
||||
from dbgpt.rag.knowledge.base import Knowledge
|
||||
from dbgpt.storage.full_text.base import FullTextStoreBase
|
||||
from dbgpt.storage.knowledge_graph.base import KnowledgeGraphBase
|
||||
from dbgpt.storage.vector_store.base import VectorStoreBase
|
||||
from dbgpt_ext.rag import ChunkParameters
|
||||
from dbgpt_ext.rag.chunk_manager import ChunkManager
|
||||
|
||||
from .base import DomainKnowledgeIndex
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DomainGeneralIndex(DomainKnowledgeIndex):
|
||||
"""General domain knowledge index for local documents."""
|
||||
|
||||
async def extract(
|
||||
self,
|
||||
knowledge: Knowledge,
|
||||
chunk_parameter: ChunkParameters,
|
||||
extract_image: bool = False,
|
||||
**kwargs,
|
||||
) -> List[Chunk]:
|
||||
if not knowledge:
|
||||
raise ValueError("knowledge must be provided.")
|
||||
documents = await knowledge.aload()
|
||||
chunk_manager = ChunkManager(
|
||||
knowledge=knowledge, chunk_parameter=chunk_parameter
|
||||
)
|
||||
chunks = chunk_manager.split(documents)
|
||||
for chunk in chunks:
|
||||
chunk.metadata["chunk_id"] = chunk.chunk_id
|
||||
if chunk_parameter.need_index_headers:
|
||||
new_chunks = []
|
||||
for chunk in chunks:
|
||||
for key, value in chunk.metadata.items():
|
||||
if value in chunk_parameter.need_index_headers:
|
||||
new_chunks.append(chunk)
|
||||
break
|
||||
return new_chunks
|
||||
if extract_image:
|
||||
new_chunks = knowledge.extract_images(chunks)
|
||||
return new_chunks
|
||||
return chunks
|
||||
|
||||
async def transform(
|
||||
self,
|
||||
chunks: List[Chunk],
|
||||
image_extractor=None,
|
||||
summary_extractor=None,
|
||||
batch_size: int = 1,
|
||||
**kwargs,
|
||||
) -> List[Chunk]:
|
||||
transform_chunks = chunks
|
||||
if image_extractor:
|
||||
transform_chunks = await self._process_images(
|
||||
chunks, image_extractor, batch_size
|
||||
)
|
||||
if summary_extractor:
|
||||
for chunk in transform_chunks:
|
||||
summary_text = await summary_extractor.extract(text=chunk.content)
|
||||
chunk.summary = summary_text
|
||||
return transform_chunks
|
||||
|
||||
async def _process_images(self, chunks, image_extractor, batch_size=1):
|
||||
import asyncio
|
||||
processed_chunks = list(chunks)
|
||||
for i in range(0, len(chunks), batch_size):
|
||||
batch = chunks[i : i + batch_size]
|
||||
tasks = []
|
||||
for chunk in batch:
|
||||
if chunk.image_url:
|
||||
tasks.append(
|
||||
self._extract_image_task(chunk, image_extractor)
|
||||
)
|
||||
if tasks:
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
for result in results:
|
||||
if result is not None and not isinstance(result, Exception):
|
||||
processed_chunks.append(result)
|
||||
return processed_chunks
|
||||
|
||||
async def _extract_image_task(self, chunk, image_extractor):
|
||||
try:
|
||||
image_text = await image_extractor.extract(
|
||||
image=chunk.image_url, text=chunk.content
|
||||
)
|
||||
chunk.content = image_text
|
||||
return chunk
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing image {chunk.image_url}: {e}")
|
||||
return None
|
||||
|
||||
async def load(
|
||||
self,
|
||||
chunks: list[Chunk],
|
||||
vector_store: Optional[VectorStoreBase] = None,
|
||||
full_text_store: Optional[FullTextStoreBase] = None,
|
||||
kg_store: Optional[KnowledgeGraphBase] = None,
|
||||
keywords: bool = True,
|
||||
max_chunks_once_load: int = 10,
|
||||
max_threads: int = 1,
|
||||
**kwargs,
|
||||
) -> List[Chunk]:
|
||||
if vector_store:
|
||||
vector_ids = await vector_store.aload_document_with_limit(
|
||||
chunks, max_chunks_once_load, max_threads
|
||||
)
|
||||
for chunk, vector_id in zip(chunks, vector_ids):
|
||||
chunk.vector_id = vector_id
|
||||
if full_text_store:
|
||||
await full_text_store.aload_document_with_limit(
|
||||
chunks, max_chunks_once_load, max_threads
|
||||
)
|
||||
if kg_store:
|
||||
await kg_store.aload_document_with_limit(
|
||||
chunks, max_chunks_once_load, max_threads
|
||||
)
|
||||
return chunks
|
||||
|
||||
async def clean(
|
||||
self,
|
||||
chunks: list[Chunk],
|
||||
node_ids: Optional[list[str]],
|
||||
with_keywords: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def domain_type(cls) -> str:
|
||||
return "normal"
|
||||
@@ -1,33 +0,0 @@
|
||||
"""Tests for DomainKnowledgeIndexFactory."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ..base import DomainKnowledgeIndex
|
||||
from ..factory import DomainKnowledgeIndexFactory
|
||||
from ..index import DomainGeneralIndex
|
||||
|
||||
|
||||
class TestDomainKnowledgeIndexFactory:
|
||||
def test_create_normal_index(self):
|
||||
index = DomainKnowledgeIndexFactory.create("normal")
|
||||
assert isinstance(index, DomainGeneralIndex)
|
||||
assert index.domain_type() == "normal"
|
||||
|
||||
def test_create_normal_index_case_insensitive(self):
|
||||
index = DomainKnowledgeIndexFactory.create("Normal")
|
||||
assert isinstance(index, DomainGeneralIndex)
|
||||
assert index.domain_type() == "normal"
|
||||
|
||||
def test_create_unknown_index_raises(self):
|
||||
with pytest.raises(Exception, match="not supported"):
|
||||
DomainKnowledgeIndexFactory.create("unknown_type")
|
||||
|
||||
def test_available_types_includes_normal(self):
|
||||
types = DomainKnowledgeIndexFactory.available_types()
|
||||
assert "normal" in types
|
||||
|
||||
def test_domain_general_index_is_subclass(self):
|
||||
assert issubclass(DomainGeneralIndex, DomainKnowledgeIndex)
|
||||
|
||||
def test_domain_general_index_domain_type(self):
|
||||
assert DomainGeneralIndex.domain_type() == "normal"
|
||||
@@ -1,163 +0,0 @@
|
||||
"""CodeGraph Build Service - Build code knowledge graphs for git repo knowledge spaces.
|
||||
|
||||
Provides functions to build code graphs from repository files and persist them
|
||||
for later querying by codegraph tools.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def build_code_graph_from_knowledge_space(
|
||||
knowledge_id: str,
|
||||
) -> Optional[Dict]:
|
||||
"""Build a code graph by reconstructing files from chunks in the DB.
|
||||
|
||||
This is a fallback method when the original files are not available
|
||||
(e.g., after a service restart). It reconstructs file contents from
|
||||
the stored chunks and builds the graph from those.
|
||||
|
||||
Args:
|
||||
knowledge_id: Knowledge space ID.
|
||||
|
||||
Returns:
|
||||
Dict with graph stats, or None if building fails.
|
||||
"""
|
||||
try:
|
||||
from ..models.chunk_db import DocumentChunkDao, DocumentChunkEntity
|
||||
from ..models.document_db import KnowledgeDocumentDao, KnowledgeDocumentEntity
|
||||
from ..tools.codegraph_tools import _save_graph
|
||||
|
||||
# Get all documents for this knowledge space
|
||||
doc_dao = KnowledgeDocumentDao()
|
||||
docs = doc_dao.get_knowledge_documents(
|
||||
KnowledgeDocumentEntity(knowledge_id=knowledge_id),
|
||||
page=1,
|
||||
page_size=100000,
|
||||
)
|
||||
if not docs:
|
||||
logger.warning(f"No documents found for knowledge_id={knowledge_id}")
|
||||
return None
|
||||
|
||||
# Reconstruct files from chunks
|
||||
import json
|
||||
|
||||
files_by_path = {}
|
||||
chunk_dao = DocumentChunkDao()
|
||||
|
||||
for doc in docs:
|
||||
try:
|
||||
meta = json.loads(doc.meta_data) if doc.meta_data else {}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
meta = {}
|
||||
|
||||
file_path = meta.get("file_path", "")
|
||||
file_type = meta.get("file_type", "")
|
||||
|
||||
# Only process code and markdown files
|
||||
if file_type not in ("code", "markdown") or not file_path:
|
||||
continue
|
||||
|
||||
# Get chunks for this document
|
||||
chunks = chunk_dao.get_document_chunks(
|
||||
DocumentChunkEntity(doc_id=doc.doc_id),
|
||||
page=1,
|
||||
page_size=10000,
|
||||
)
|
||||
|
||||
# Reconstruct file content from chunks
|
||||
content_parts = []
|
||||
for chunk in sorted(chunks, key=lambda c: c.id if hasattr(c, 'id') else 0):
|
||||
if chunk.chunk_type == "summary":
|
||||
continue
|
||||
content_parts.append(chunk.content or "")
|
||||
|
||||
full_content = "\n".join(content_parts)
|
||||
if full_content.strip():
|
||||
files_by_path[file_path] = full_content
|
||||
|
||||
if not files_by_path:
|
||||
logger.warning(f"No code files found for knowledge_id={knowledge_id}")
|
||||
return None
|
||||
|
||||
# Build graph from reconstructed files
|
||||
files_list = [
|
||||
{"path": path, "content": content}
|
||||
for path, content in files_by_path.items()
|
||||
]
|
||||
|
||||
from dbgpt_ext.rag.graph_builder.repo_graph_builder import RepoGraphBuilder
|
||||
|
||||
builder = RepoGraphBuilder()
|
||||
graph = await builder.build_from_files(
|
||||
files=files_list,
|
||||
repo_name=knowledge_id,
|
||||
)
|
||||
|
||||
if graph and graph.vertex_count > 0:
|
||||
_save_graph(knowledge_id, graph, build_source="chunk_reconstruction")
|
||||
return {
|
||||
"vertices": graph.vertex_count,
|
||||
"edges": graph.edge_count,
|
||||
"files_processed": len(files_list),
|
||||
"status": "completed",
|
||||
}
|
||||
else:
|
||||
logger.warning(f"Graph building produced empty graph for {knowledge_id}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to build code graph from chunks: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def build_code_graph_from_files(
|
||||
knowledge_id: str,
|
||||
files: List[Dict],
|
||||
repo_url: str = "",
|
||||
repo_name: str = "",
|
||||
) -> Optional[Dict]:
|
||||
"""Build a code graph from a list of file dicts.
|
||||
|
||||
Args:
|
||||
knowledge_id: Knowledge space ID.
|
||||
files: List of dicts with 'path' and 'content' keys.
|
||||
repo_url: Repository URL.
|
||||
repo_name: Repository name.
|
||||
|
||||
Returns:
|
||||
Dict with graph stats, or None if building fails.
|
||||
"""
|
||||
try:
|
||||
from ..tools.codegraph_tools import _save_graph
|
||||
from dbgpt_ext.rag.graph_builder.repo_graph_builder import RepoGraphBuilder
|
||||
|
||||
builder = RepoGraphBuilder()
|
||||
graph = await builder.build_from_files(
|
||||
files=files,
|
||||
repo_url=repo_url,
|
||||
repo_name=repo_name or knowledge_id,
|
||||
)
|
||||
|
||||
if graph and graph.vertex_count > 0:
|
||||
_save_graph(knowledge_id, graph, build_source="files")
|
||||
return {
|
||||
"vertices": graph.vertex_count,
|
||||
"edges": graph.edge_count,
|
||||
"files_processed": len(files),
|
||||
"status": "completed",
|
||||
}
|
||||
else:
|
||||
logger.warning(
|
||||
f"Graph building produced empty graph for {knowledge_id}"
|
||||
)
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to build code graph from files: {e}")
|
||||
return None
|
||||
@@ -1,22 +0,0 @@
|
||||
"""Knowledge base search tools package.
|
||||
|
||||
Provides structured search tools for agents:
|
||||
- kb_ls: List files and directories in a knowledge space
|
||||
- kb_glob: Search files by name pattern
|
||||
- kb_grep: Search file contents by keyword
|
||||
- kb_cat: Read file content by path
|
||||
- kb_semantic_search: Semantic search using vector retrieval
|
||||
- kb_codegraph_explore: Code knowledge graph exploration
|
||||
"""
|
||||
|
||||
# Import tool modules to register them with the @tool decorator
|
||||
from . import kb_file_tools # noqa: F401
|
||||
from . import semantic_search_tool # noqa: F401
|
||||
|
||||
# CodeGraph tools - optional, requires graph_store
|
||||
try:
|
||||
from . import codegraph_tools # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
__all__ = ["kb_file_tools", "semantic_search_tool", "codegraph_tools"]
|
||||
@@ -1,21 +0,0 @@
|
||||
"""Code graph tools for GIT_REPO knowledge spaces."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Annotated
|
||||
from dbgpt.agent.resource.tool.base import tool
|
||||
from dbgpt.storage.graph_store.graph import MemoryGraph
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_graph_cache: dict = {}
|
||||
|
||||
|
||||
def _get_graph_cache_dir(knowledge_id):
|
||||
return os.path.join(os.path.expanduser("~"), ".dbgpt", "graph_cache", knowledge_id)
|
||||
|
||||
|
||||
def _load_graph(knowledge_id):
|
||||
if knowledge_id in _graph_cache:
|
||||
return _graph_cache[knowledge_id], None
|
||||
graph_file = os.path.join(_get_graph_cache_dir(knowledge_id
|
||||
@@ -1,344 +0,0 @@
|
||||
"""File system tools for knowledge spaces (kb_ls, kb_glob, kb_grep, kb_cat)."""
|
||||
|
||||
import fnmatch
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Annotated, List, Optional
|
||||
|
||||
from dbgpt.agent.resource.tool.base import tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_CONTEXT_PREFIX_RE = re.compile(r"^\[文件: [^\]]*\](?:\s*\[路径: [^\]]*\])?\s*\n?")
|
||||
|
||||
|
||||
def _get_document_dao():
|
||||
from ..models.document_db import KnowledgeDocumentDao, KnowledgeDocumentEntity
|
||||
return KnowledgeDocumentDao(), KnowledgeDocumentEntity
|
||||
|
||||
|
||||
def _get_chunk_dao():
|
||||
from ..models.chunk_db import DocumentChunkDao, DocumentChunkEntity
|
||||
return DocumentChunkDao(), DocumentChunkEntity
|
||||
|
||||
|
||||
def _get_all_file_paths(knowledge_id: str) -> List[dict]:
|
||||
dao, Entity = _get_document_dao()
|
||||
docs = dao.get_knowledge_documents(
|
||||
Entity(knowledge_id=knowledge_id), page=1, page_size=10000,
|
||||
)
|
||||
if not docs:
|
||||
return []
|
||||
results = []
|
||||
for doc in docs:
|
||||
try:
|
||||
meta = json.loads(doc.meta_data) if doc.meta_data else {}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
meta = {}
|
||||
if not isinstance(meta, dict):
|
||||
meta = {}
|
||||
file_path = meta.get("file_path", "") or doc.doc_name or ""
|
||||
if not file_path:
|
||||
continue
|
||||
results.append({
|
||||
"file_path": file_path,
|
||||
"file_type": meta.get("file_type", ""),
|
||||
"language": meta.get("language", ""),
|
||||
"doc_id": doc.doc_id,
|
||||
"doc_name": doc.doc_name,
|
||||
})
|
||||
return results
|
||||
|
||||
|
||||
def _find_doc_by_file_path(knowledge_id: str, path: str) -> Optional[dict]:
|
||||
all_files = _get_all_file_paths(knowledge_id)
|
||||
for f in all_files:
|
||||
if f["file_path"] == path:
|
||||
return f
|
||||
return None
|
||||
|
||||
|
||||
@tool(
|
||||
name="kb_ls",
|
||||
description="列出知识库指定目录下的文件和子目录。查找文件请优先用 kb_glob 或 kb_grep。",
|
||||
)
|
||||
async def kb_ls(
|
||||
knowledge_id: Annotated[str, "知识空间 ID"],
|
||||
path: Annotated[str, "目录路径，空字符串表示根目录"] = "",
|
||||
offset: Annotated[int, "跳过前 N 条记录"] = 0,
|
||||
limit: Annotated[int, "最多返回条目数"] = 200,
|
||||
) -> str:
|
||||
offset = int(offset) if offset else 0
|
||||
limit = int(limit) if limit else 200
|
||||
all_files = _get_all_file_paths(knowledge_id)
|
||||
if not all_files:
|
||||
return f"知识空间 {knowledge_id} 中没有文件"
|
||||
target = path.rstrip("/")
|
||||
prefix = (target + "/") if target else ""
|
||||
dirs = {}
|
||||
files = []
|
||||
for f in all_files:
|
||||
fp = f["file_path"]
|
||||
if not fp.startswith(prefix):
|
||||
continue
|
||||
remaining = fp[len(prefix):]
|
||||
if not remaining:
|
||||
continue
|
||||
parts = remaining.split("/")
|
||||
if len(parts) == 1:
|
||||
files.append((parts[0], f["file_type"], f["language"]))
|
||||
else:
|
||||
dirs[parts[0]] = dirs.get(parts[0], 0) + 1
|
||||
if not dirs and not files:
|
||||
return f"目录 '{path}' 不存在或为空"
|
||||
all_entries = []
|
||||
for dir_name in sorted(dirs.keys()):
|
||||
all_entries.append(f" {dir_name}/\t({dirs[dir_name]} files)")
|
||||
for name, ftype, lang in sorted(files):
|
||||
all_entries.append(f" {name}\t{lang or ftype or ''}")
|
||||
total_entries = len(all_entries)
|
||||
display_path = target or "/"
|
||||
paged_entries = all_entries[offset : offset + limit]
|
||||
if not paged_entries:
|
||||
return f"目录 '{display_path}' 共 {total_entries} 条，offset={offset} 超出范围"
|
||||
lines = [f"目录: {display_path} ({len(files)} files, {len(dirs)} dirs)"]
|
||||
MAX_OUTPUT_CHARS = 8000
|
||||
current_chars = len(lines[0])
|
||||
truncated = False
|
||||
for entry in paged_entries:
|
||||
current_chars += len(entry)
|
||||
if current_chars > MAX_OUTPUT_CHARS:
|
||||
truncated = True
|
||||
break
|
||||
lines.append(entry)
|
||||
shown = len(lines) - 1
|
||||
if truncated or offset + limit < total_entries:
|
||||
next_offset = offset + shown
|
||||
remaining = total_entries - next_offset
|
||||
if remaining > 0:
|
||||
lines.append(f"\n... 还有 {remaining} 条未显示，使用 offset={next_offset} 查看后续结果。")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@tool(
|
||||
name="kb_glob",
|
||||
description="按文件名/文档名搜索知识库文件。支持关键词和 glob 模式。",
|
||||
)
|
||||
async def kb_glob(
|
||||
knowledge_id: Annotated[str, "知识空间 ID"],
|
||||
pattern: Annotated[str, "文件名关键词或 glob 模式"],
|
||||
limit: Annotated[int, "最多返回文件数"] = 200,
|
||||
offset: Annotated[int, "跳过前 N 个匹配文件"] = 0,
|
||||
) -> str:
|
||||
limit = int(limit) if limit else 200
|
||||
offset = int(offset) if offset else 0
|
||||
all_files = _get_all_file_paths(knowledge_id)
|
||||
if not all_files:
|
||||
return f"知识空间 {knowledge_id} 中没有文件"
|
||||
is_glob = any(c in pattern for c in "*?[")
|
||||
matches = []
|
||||
for f in all_files:
|
||||
fp = f["file_path"]
|
||||
if is_glob:
|
||||
if fnmatch.fnmatch(fp, pattern):
|
||||
matches.append(f)
|
||||
elif pattern.startswith("**/") and fnmatch.fnmatch(fp, pattern[3:]):
|
||||
matches.append(f)
|
||||
else:
|
||||
if pattern.lower() in fp.lower():
|
||||
matches.append(f)
|
||||
if not matches:
|
||||
return f"没有匹配 '{pattern}' 的文件"
|
||||
total_matches = len(matches)
|
||||
sorted_matches = sorted(matches, key=lambda x: x["file_path"])
|
||||
paged = sorted_matches[offset : offset + limit]
|
||||
if not paged:
|
||||
return f"匹配 '{pattern}' 共 {total_matches} 个文件，offset={offset} 超出范围"
|
||||
lines = [f"匹配 '{pattern}' 共 {total_matches} 个文件 (显示第 {offset + 1}-{offset + len(paged)} 个):"]
|
||||
MAX_OUTPUT_CHARS = 8000
|
||||
current_chars = len(lines[0])
|
||||
truncated = False
|
||||
for f in paged:
|
||||
type_info = f["language"] or f["file_type"] or ""
|
||||
entry = f" {f['file_path']}\t{type_info}"
|
||||
current_chars += len(entry)
|
||||
if current_chars > MAX_OUTPUT_CHARS:
|
||||
truncated = True
|
||||
break
|
||||
lines.append(entry)
|
||||
shown = len(lines) - 1
|
||||
if truncated or offset + limit < total_matches:
|
||||
next_offset = offset + shown
|
||||
remaining = total_matches - next_offset
|
||||
if remaining > 0:
|
||||
lines.append(f"\n... 还有 {remaining} 个文件未显示，使用 offset={next_offset} 查看后续结果。")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@tool(
|
||||
name="kb_grep",
|
||||
description="在知识库文件内容中搜索关键词，返回包含该关键词的文件和具体行内容。优先使用此工具而非语义搜索。",
|
||||
)
|
||||
async def kb_grep(
|
||||
knowledge_id: Annotated[str, "知识空间 ID"],
|
||||
query: Annotated[str, "搜索关键词"],
|
||||
path: Annotated[str, "限定搜索的目录路径"] = "",
|
||||
file_pattern: Annotated[str, "限定文件类型如 '*.py'"] = "",
|
||||
limit: Annotated[int, "最多返回文件数"] = 20,
|
||||
offset: Annotated[int, "跳过前 N 个匹配文件"] = 0,
|
||||
) -> str:
|
||||
limit = int(limit) if limit else 20
|
||||
offset = int(offset) if offset else 0
|
||||
chunk_dao, ChunkEntity = _get_chunk_dao()
|
||||
all_files = _get_all_file_paths(knowledge_id)
|
||||
if not all_files:
|
||||
return f"知识空间 {knowledge_id} 中没有文件"
|
||||
target_files = {}
|
||||
norm_path = path.rstrip("/") if path else ""
|
||||
for f in all_files:
|
||||
fp = f["file_path"]
|
||||
if norm_path:
|
||||
if fp != norm_path and not fp.startswith(norm_path + "/"):
|
||||
continue
|
||||
if file_pattern and not fnmatch.fnmatch(fp, file_pattern):
|
||||
continue
|
||||
target_files[f["doc_id"]] = f
|
||||
if not target_files:
|
||||
scope = path or file_pattern or "仓库"
|
||||
return f"在 '{scope}' 范围内没有文件"
|
||||
doc_matches = {}
|
||||
fetch_limit = offset + limit
|
||||
if len(target_files) <= 50:
|
||||
for doc_id in target_files:
|
||||
if len(doc_matches) >= fetch_limit:
|
||||
break
|
||||
chunks = chunk_dao.get_document_chunks(
|
||||
ChunkEntity(doc_id=doc_id, content=query), page=1, page_size=200,
|
||||
)
|
||||
_collect_grep_matches(chunks, doc_id, query, doc_matches, fetch_limit)
|
||||
else:
|
||||
query_entity = ChunkEntity(knowledge_id=knowledge_id, content=query)
|
||||
chunks = chunk_dao.get_document_chunks(query_entity, page=1, page_size=500)
|
||||
for chunk in chunks:
|
||||
if chunk.chunk_type == "summary":
|
||||
continue
|
||||
doc_id = chunk.doc_id
|
||||
if doc_id not in target_files:
|
||||
continue
|
||||
if len(doc_matches) >= fetch_limit and doc_id not in doc_matches:
|
||||
break
|
||||
_collect_grep_matches([chunk], doc_id, query, doc_matches, fetch_limit)
|
||||
if not doc_matches:
|
||||
scope = path or file_pattern or "仓库"
|
||||
return f"在 '{scope}' 中未找到包含 '{query}' 的内容"
|
||||
all_doc_ids = list(doc_matches.keys())
|
||||
total_files = len(all_doc_ids)
|
||||
total_matches = sum(len(v) for v in doc_matches.values())
|
||||
paged_doc_ids = all_doc_ids[offset : offset + limit]
|
||||
if not paged_doc_ids:
|
||||
return f"'{query}' 共匹配 {total_files} 个文件，offset={offset} 超出范围"
|
||||
result_lines = [
|
||||
f"'{query}' 匹配 {total_files} 个文件 {total_matches} 处"
|
||||
f" (显示第 {offset + 1}-{offset + len(paged_doc_ids)} 个文件):"
|
||||
]
|
||||
MAX_OUTPUT_CHARS = 8000
|
||||
current_chars = len(result_lines[0])
|
||||
truncated = False
|
||||
for doc_id in paged_doc_ids:
|
||||
matches = doc_matches[doc_id]
|
||||
file_info = target_files.get(doc_id) or {}
|
||||
file_header = f"\n{file_info.get('file_path', doc_id)}:"
|
||||
result_lines.append(file_header)
|
||||
current_chars += len(file_header)
|
||||
for line_no, content in matches[:10]:
|
||||
line = f" {line_no}: {content}"
|
||||
current_chars += len(line)
|
||||
if current_chars > MAX_OUTPUT_CHARS:
|
||||
truncated = True
|
||||
break
|
||||
result_lines.append(line)
|
||||
if truncated:
|
||||
break
|
||||
if truncated or offset + limit < total_files:
|
||||
next_offset = offset + len(paged_doc_ids)
|
||||
result_lines.append(f"\n... 使用 offset={next_offset} 查看后续结果。")
|
||||
return "\n".join(result_lines)
|
||||
|
||||
|
||||
def _collect_grep_matches(chunks, doc_id, query, doc_matches, limit):
|
||||
for chunk in chunks:
|
||||
if chunk.chunk_type == "summary":
|
||||
continue
|
||||
if doc_id not in doc_matches:
|
||||
if len(doc_matches) >= limit:
|
||||
return
|
||||
doc_matches[doc_id] = []
|
||||
content = chunk.content or ""
|
||||
content = _CONTEXT_PREFIX_RE.sub("", content)
|
||||
try:
|
||||
chunk_meta = json.loads(chunk.meta_data) if chunk.meta_data else {}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
chunk_meta = {}
|
||||
start_line = chunk_meta.get("start_line", 1)
|
||||
lines = content.split("\n")
|
||||
for i, line in enumerate(lines):
|
||||
if query.lower() in line.lower():
|
||||
line_no = start_line + i
|
||||
doc_matches[doc_id].append((line_no, line.strip()[:120]))
|
||||
|
||||
|
||||
@tool(
|
||||
name="kb_cat",
|
||||
description="读取知识库中指定文件/文档的内容。支持按行号范围读取。",
|
||||
)
|
||||
async def kb_cat(
|
||||
knowledge_id: Annotated[str, "知识空间 ID"],
|
||||
path: Annotated[str, "文件路径，如 'src/auth/login.py'"],
|
||||
start_line: Annotated[int, "起始行号（从 1 开始）"] = 1,
|
||||
end_line: Annotated[int, "结束行号，0 表示读到末尾"] = 0,
|
||||
) -> str:
|
||||
start_line = int(start_line) if start_line else 1
|
||||
end_line = int(end_line) if end_line else 0
|
||||
file_info = _find_doc_by_file_path(knowledge_id, path)
|
||||
if not file_info:
|
||||
return f"文件 '{path}' 不存在"
|
||||
doc_id = file_info["doc_id"]
|
||||
chunk_dao, ChunkEntity = _get_chunk_dao()
|
||||
chunks = chunk_dao.get_document_chunks(
|
||||
ChunkEntity(doc_id=doc_id), page=1, page_size=1000,
|
||||
)
|
||||
content_chunks = []
|
||||
for chunk in chunks:
|
||||
if chunk.chunk_type == "summary":
|
||||
continue
|
||||
try:
|
||||
meta = json.loads(chunk.meta_data) if chunk.meta_data else {}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
meta = {}
|
||||
chunk_index = meta.get("chunk_index", 0)
|
||||
content_chunks.append((chunk_index, chunk.content or ""))
|
||||
content_chunks.sort(key=lambda x: x[0])
|
||||
full_lines = []
|
||||
for _, content in content_chunks:
|
||||
content = _CONTEXT_PREFIX_RE.sub("", content)
|
||||
full_lines.extend(content.split("\n"))
|
||||
if not full_lines:
|
||||
return f"文件 '{path}' 内容为空"
|
||||
total_lines = len(full_lines)
|
||||
lang = file_info.get("language") or file_info.get("file_type") or ""
|
||||
start_idx = max(0, start_line - 1)
|
||||
end_idx = end_line if end_line > 0 else total_lines
|
||||
end_idx = min(end_idx, total_lines)
|
||||
selected = full_lines[start_idx:end_idx]
|
||||
result_lines = [f"{path} ({lang}, {total_lines} lines)"]
|
||||
for i, line in enumerate(selected):
|
||||
line_no = start_idx + i + 1
|
||||
result_lines.append(f" {line_no:>4} | {line}")
|
||||
if len(result_lines) > 502:
|
||||
result_lines = result_lines[:502]
|
||||
next_start = start_idx + 500 + 1
|
||||
result_lines.append(
|
||||
f" ... (已截断，使用 start_line={next_start} 继续读取)"
|
||||
)
|
||||
return "\n".join(result_lines)
|
||||
@@ -1,84 +0,0 @@
|
||||
"""Semantic search tool for knowledge spaces."""
|
||||
|
||||
import logging
|
||||
from typing import Annotated
|
||||
|
||||
from dbgpt.agent.resource.tool.base import tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _get_rag_service():
|
||||
from ..service.service import Service
|
||||
from dbgpt._private.config import Config
|
||||
system_app = Config().SYSTEM_APP
|
||||
if not system_app:
|
||||
raise RuntimeError("SYSTEM_APP is not initialized yet")
|
||||
return Service.get_instance(system_app)
|
||||
|
||||
|
||||
def _format_chunk_results(chunks, query: str, knowledge_id: str) -> str:
|
||||
result_lines = [f"Semantic search '{query}' in knowledge space {knowledge_id}:"]
|
||||
MAX_OUTPUT_CHARS = 8000
|
||||
current_chars = len(result_lines[0])
|
||||
for i, chunk in enumerate(chunks):
|
||||
content = chunk.content or ""
|
||||
score = getattr(chunk, "score", None)
|
||||
metadata = getattr(chunk, "metadata", {}) or {}
|
||||
file_path = metadata.get("file_path", "")
|
||||
doc_name = metadata.get("doc_name", "")
|
||||
score_str = f" (score: {score:.2f})" if score is not None else ""
|
||||
source_parts = []
|
||||
if file_path:
|
||||
source_parts.append(f"file: {file_path}")
|
||||
if doc_name and doc_name != file_path:
|
||||
source_parts.append(f"doc: {doc_name}")
|
||||
source_str = " | " + " | ".join(source_parts) if source_parts else ""
|
||||
entry = f"\n---\n### Result {i + 1}{score_str}{source_str}\n{content}"
|
||||
current_chars += len(entry)
|
||||
if current_chars > MAX_OUTPUT_CHARS:
|
||||
result_lines.append(f"\n... {len(chunks) - i} more results not shown.")
|
||||
break
|
||||
result_lines.append(entry)
|
||||
return "\n".join(result_lines)
|
||||
|
||||
|
||||
@tool(
|
||||
name="kb_semantic_search",
|
||||
description=(
|
||||
"Semantic search in knowledge base. "
|
||||
"Priority: kb_grep (exact match) > kb_semantic_search (semantic match). "
|
||||
"Use this only when kb_grep returns empty or insufficient results."
|
||||
),
|
||||
)
|
||||
async def kb_semantic_search(
|
||||
knowledge_id: Annotated[str, "Knowledge space ID"],
|
||||
query: Annotated[str, "Search query in natural language"],
|
||||
top_k: Annotated[int, "Number of results"] = 5,
|
||||
score_threshold: Annotated[float, "Minimum score threshold (0-1)"] = 0.0,
|
||||
) -> str:
|
||||
top_k = int(top_k) if top_k else 5
|
||||
score_threshold = float(score_threshold) if score_threshold else 0.0
|
||||
try:
|
||||
service = _get_rag_service()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get RAG service: {e}")
|
||||
return f"Semantic search service unavailable: {e}"
|
||||
try:
|
||||
from ..api.schemas import KnowledgeRetrieveRequest
|
||||
request = KnowledgeRetrieveRequest(
|
||||
query=query,
|
||||
space_id=int(knowledge_id) if knowledge_id.isdigit() else knowledge_id,
|
||||
top_k=top_k,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
space = service.get({"id": request.space_id})
|
||||
if space is None:
|
||||
return f"Knowledge space {knowledge_id} not found"
|
||||
search_res = await service.retrieve(request, space)
|
||||
except Exception as e:
|
||||
logger.exception(f"Semantic search failed: {e}")
|
||||
return f"Semantic search failed: {e}"
|
||||
if not search_res:
|
||||
return f"No results found for '{query}' in knowledge space {knowledge_id}"
|
||||
return _format_chunk_results(search_res, query, knowledge_id)
|
||||
@@ -1,64 +0,0 @@
|
||||
"""Semantic search tool for knowledge spaces."""
|
||||
|
||||
import logging
|
||||
from typing import Annotated
|
||||
from dbgpt.agent.resource.tool.base import tool
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _get_rag_service():
|
||||
from ..service.service import Service
|
||||
from dbgpt._private.config import Config
|
||||
system_app = Config().SYSTEM_APP
|
||||
if not system_app:
|
||||
raise RuntimeError("SYSTEM_APP is not initialized yet")
|
||||
return Service.get_instance(system_app)
|
||||
|
||||
|
||||
def _format_chunk_results(chunks, query, knowledge_id):
|
||||
lines = [f"Semantic search '{query}' in {knowledge_id}:"]
|
||||
chars = len(lines[0])
|
||||
for i, chunk in enumerate(chunks):
|
||||
content = chunk.content or ""
|
||||
score = getattr(chunk, "score", None)
|
||||
metadata = getattr(chunk, "metadata", {}) or {}
|
||||
file_path = metadata.get("file_path", "")
|
||||
score_str = f" (score: {score:.2f})" if score is not None else ""
|
||||
entry = f"\n---\n### Result {i+1}{score_str} [{file_path}]\n{content}"
|
||||
chars += len(entry)
|
||||
if chars > 8000:
|
||||
lines.append(f"\n... {len(chunks)-i} more results.")
|
||||
break
|
||||
lines.append(entry)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@tool(name="kb_semantic_search", description="Semantic search in knowledge base. Use only when kb_grep returns insufficient results.")
|
||||
async def kb_semantic_search(
|
||||
knowledge_id: Annotated[str, "Knowledge space ID"],
|
||||
query: Annotated[str, "Natural language search query"],
|
||||
top_k: Annotated[int, "Number of results"] = 5,
|
||||
score_threshold: Annotated[float, "Min score (0-1)"] = 0.0,
|
||||
) -> str:
|
||||
top_k = int(top_k) if top_k else 5
|
||||
score_threshold = float(score_threshold) if score_threshold else 0.0
|
||||
try:
|
||||
service = _get_rag_service()
|
||||
except Exception as e:
|
||||
return f"Semantic search service unavailable: {e}"
|
||||
try:
|
||||
from ..api.schemas import KnowledgeRetrieveRequest
|
||||
request = KnowledgeRetrieveRequest(
|
||||
query=query,
|
||||
space_id=int(knowledge_id) if knowledge_id.isdigit() else knowledge_id,
|
||||
top_k=top_k, score_threshold=score_threshold)
|
||||
space = service.get({"id": request.space_id})
|
||||
if space is None:
|
||||
return f"Knowledge space {knowledge_id} not found"
|
||||
search_res = await service.retrieve(request, space)
|
||||
except Exception as e:
|
||||
return f"Semantic search failed: {e}"
|
||||
if not search_res:
|
||||
return f"No results for '{query}' in {knowledge_id}"
|
||||
return _format_chunk_results(search_res, query, knowledge_id)
|
||||
Reference in New Issue
Block a user