294 lines
10 KiB
Python
294 lines
10 KiB
Python
|
"""
|
|||
|
|
RAG 检索工具
|
|||
|
|
支持语义检索代码
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from typing import Optional, List
|
|||
|
|
from pydantic import BaseModel, Field
|
|||
|
|
|
|||
|
|
from .base import AgentTool, ToolResult
|
|||
|
|
from app.services.rag import CodeRetriever
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RAGQueryInput(BaseModel):
|
|||
|
|
"""RAG 查询输入参数"""
|
|||
|
|
query: str = Field(description="搜索查询，描述你要找的代码功能或特征")
|
|||
|
|
top_k: int = Field(default=10, description="返回结果数量")
|
|||
|
|
file_path: Optional[str] = Field(default=None, description="限定搜索的文件路径")
|
|||
|
|
language: Optional[str] = Field(default=None, description="限定编程语言")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RAGQueryTool(AgentTool):
|
|||
|
|
"""
|
|||
|
|
RAG 代码检索工具
|
|||
|
|
使用语义搜索在代码库中查找相关代码
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, retriever: CodeRetriever):
|
|||
|
|
super().__init__()
|
|||
|
|
self.retriever = retriever
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def name(self) -> str:
|
|||
|
|
return "rag_query"
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def description(self) -> str:
|
|||
|
|
return """在代码库中进行语义搜索。
|
|||
|
|
使用场景:
|
|||
|
|
- 查找特定功能的实现代码
|
|||
|
|
- 查找调用某个函数的代码
|
|||
|
|
- 查找处理用户输入的代码
|
|||
|
|
- 查找数据库操作相关代码
|
|||
|
|
- 查找认证/授权相关代码
|
|||
|
|
|
|||
|
|
输入:
|
|||
|
|
- query: 描述你要查找的代码，例如 "处理用户登录的函数"、"SQL查询执行"、"文件上传处理"
|
|||
|
|
- top_k: 返回结果数量（默认10）
|
|||
|
|
- file_path: 可选，限定在某个文件中搜索
|
|||
|
|
- language: 可选，限定编程语言
|
|||
|
|
|
|||
|
|
输出: 相关的代码片段列表，包含文件路径、行号、代码内容和相似度分数"""
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def args_schema(self):
|
|||
|
|
return RAGQueryInput
|
|||
|
|
|
|||
|
|
async def _execute(
|
|||
|
|
self,
|
|||
|
|
query: str,
|
|||
|
|
top_k: int = 10,
|
|||
|
|
file_path: Optional[str] = None,
|
|||
|
|
language: Optional[str] = None,
|
|||
|
|
**kwargs
|
|||
|
|
) -> ToolResult:
|
|||
|
|
"""执行 RAG 检索"""
|
|||
|
|
try:
|
|||
|
|
results = await self.retriever.retrieve(
|
|||
|
|
query=query,
|
|||
|
|
top_k=top_k,
|
|||
|
|
filter_file_path=file_path,
|
|||
|
|
filter_language=language,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not results:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
data="没有找到相关代码",
|
|||
|
|
metadata={"query": query, "results_count": 0}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 格式化输出
|
|||
|
|
output_parts = [f"找到 {len(results)} 个相关代码片段:\n"]
|
|||
|
|
|
|||
|
|
for i, result in enumerate(results):
|
|||
|
|
output_parts.append(f"\n--- 结果 {i+1} (相似度: {result.score:.2f}) ---")
|
|||
|
|
output_parts.append(f"文件: {result.file_path}")
|
|||
|
|
output_parts.append(f"行号: {result.line_start}-{result.line_end}")
|
|||
|
|
if result.name:
|
|||
|
|
output_parts.append(f"名称: {result.name}")
|
|||
|
|
if result.security_indicators:
|
|||
|
|
output_parts.append(f"安全指标: {', '.join(result.security_indicators)}")
|
|||
|
|
output_parts.append(f"代码:\n```{result.language}\n{result.content}\n```")
|
|||
|
|
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
data="\n".join(output_parts),
|
|||
|
|
metadata={
|
|||
|
|
"query": query,
|
|||
|
|
"results_count": len(results),
|
|||
|
|
"results": [r.to_dict() for r in results],
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error=f"RAG 检索失败: {str(e)}",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SecurityCodeSearchInput(BaseModel):
|
|||
|
|
"""安全代码搜索输入"""
|
|||
|
|
vulnerability_type: str = Field(
|
|||
|
|
description="漏洞类型: sql_injection, xss, command_injection, path_traversal, ssrf, deserialization, auth_bypass, hardcoded_secret"
|
|||
|
|
)
|
|||
|
|
top_k: int = Field(default=20, description="返回结果数量")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SecurityCodeSearchTool(AgentTool):
|
|||
|
|
"""
|
|||
|
|
安全相关代码搜索工具
|
|||
|
|
专门用于查找可能存在安全漏洞的代码
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, retriever: CodeRetriever):
|
|||
|
|
super().__init__()
|
|||
|
|
self.retriever = retriever
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def name(self) -> str:
|
|||
|
|
return "security_code_search"
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def description(self) -> str:
|
|||
|
|
return """搜索可能存在安全漏洞的代码。
|
|||
|
|
专门针对特定漏洞类型进行搜索。
|
|||
|
|
|
|||
|
|
支持的漏洞类型:
|
|||
|
|
- sql_injection: SQL 注入
|
|||
|
|
- xss: 跨站脚本
|
|||
|
|
- command_injection: 命令注入
|
|||
|
|
- path_traversal: 路径遍历
|
|||
|
|
- ssrf: 服务端请求伪造
|
|||
|
|
- deserialization: 不安全的反序列化
|
|||
|
|
- auth_bypass: 认证绕过
|
|||
|
|
- hardcoded_secret: 硬编码密钥"""
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def args_schema(self):
|
|||
|
|
return SecurityCodeSearchInput
|
|||
|
|
|
|||
|
|
async def _execute(
|
|||
|
|
self,
|
|||
|
|
vulnerability_type: str,
|
|||
|
|
top_k: int = 20,
|
|||
|
|
**kwargs
|
|||
|
|
) -> ToolResult:
|
|||
|
|
"""执行安全代码搜索"""
|
|||
|
|
try:
|
|||
|
|
results = await self.retriever.retrieve_security_related(
|
|||
|
|
vulnerability_type=vulnerability_type,
|
|||
|
|
top_k=top_k,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not results:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
data=f"没有找到与 {vulnerability_type} 相关的代码",
|
|||
|
|
metadata={"vulnerability_type": vulnerability_type, "results_count": 0}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 格式化输出
|
|||
|
|
output_parts = [f"找到 {len(results)} 个可能与 {vulnerability_type} 相关的代码:\n"]
|
|||
|
|
|
|||
|
|
for i, result in enumerate(results):
|
|||
|
|
output_parts.append(f"\n--- 可疑代码 {i+1} ---")
|
|||
|
|
output_parts.append(f"文件: {result.file_path}:{result.line_start}")
|
|||
|
|
if result.security_indicators:
|
|||
|
|
output_parts.append(f"⚠️ 安全指标: {', '.join(result.security_indicators)}")
|
|||
|
|
output_parts.append(f"代码:\n```{result.language}\n{result.content}\n```")
|
|||
|
|
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
data="\n".join(output_parts),
|
|||
|
|
metadata={
|
|||
|
|
"vulnerability_type": vulnerability_type,
|
|||
|
|
"results_count": len(results),
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error=f"安全代码搜索失败: {str(e)}",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class FunctionContextInput(BaseModel):
|
|||
|
|
"""函数上下文搜索输入"""
|
|||
|
|
function_name: str = Field(description="函数名称")
|
|||
|
|
file_path: Optional[str] = Field(default=None, description="文件路径")
|
|||
|
|
include_callers: bool = Field(default=True, description="是否包含调用者")
|
|||
|
|
include_callees: bool = Field(default=True, description="是否包含被调用的函数")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class FunctionContextTool(AgentTool):
|
|||
|
|
"""
|
|||
|
|
函数上下文搜索工具
|
|||
|
|
查找函数的定义、调用者和被调用者
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, retriever: CodeRetriever):
|
|||
|
|
super().__init__()
|
|||
|
|
self.retriever = retriever
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def name(self) -> str:
|
|||
|
|
return "function_context"
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def description(self) -> str:
|
|||
|
|
return """查找函数的上下文信息，包括定义、调用者和被调用的函数。
|
|||
|
|
用于追踪数据流和理解函数的使用方式。
|
|||
|
|
|
|||
|
|
输入:
|
|||
|
|
- function_name: 要查找的函数名
|
|||
|
|
- file_path: 可选，限定文件路径
|
|||
|
|
- include_callers: 是否查找调用此函数的代码
|
|||
|
|
- include_callees: 是否查找此函数调用的其他函数"""
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def args_schema(self):
|
|||
|
|
return FunctionContextInput
|
|||
|
|
|
|||
|
|
async def _execute(
|
|||
|
|
self,
|
|||
|
|
function_name: str,
|
|||
|
|
file_path: Optional[str] = None,
|
|||
|
|
include_callers: bool = True,
|
|||
|
|
include_callees: bool = True,
|
|||
|
|
**kwargs
|
|||
|
|
) -> ToolResult:
|
|||
|
|
"""执行函数上下文搜索"""
|
|||
|
|
try:
|
|||
|
|
context = await self.retriever.retrieve_function_context(
|
|||
|
|
function_name=function_name,
|
|||
|
|
file_path=file_path,
|
|||
|
|
include_callers=include_callers,
|
|||
|
|
include_callees=include_callees,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
output_parts = [f"函数 '{function_name}' 的上下文分析:\n"]
|
|||
|
|
|
|||
|
|
# 函数定义
|
|||
|
|
if context["definition"]:
|
|||
|
|
output_parts.append("### 函数定义:")
|
|||
|
|
for result in context["definition"]:
|
|||
|
|
output_parts.append(f"文件: {result.file_path}:{result.line_start}")
|
|||
|
|
output_parts.append(f"```{result.language}\n{result.content}\n```")
|
|||
|
|
else:
|
|||
|
|
output_parts.append("未找到函数定义")
|
|||
|
|
|
|||
|
|
# 调用者
|
|||
|
|
if context["callers"]:
|
|||
|
|
output_parts.append(f"\n### 调用此函数的代码 ({len(context['callers'])} 处):")
|
|||
|
|
for result in context["callers"][:5]:
|
|||
|
|
output_parts.append(f"- {result.file_path}:{result.line_start}")
|
|||
|
|
output_parts.append(f"```{result.language}\n{result.content[:500]}\n```")
|
|||
|
|
|
|||
|
|
# 被调用者
|
|||
|
|
if context["callees"]:
|
|||
|
|
output_parts.append(f"\n### 此函数调用的其他函数:")
|
|||
|
|
for result in context["callees"][:5]:
|
|||
|
|
if result.name:
|
|||
|
|
output_parts.append(f"- {result.name} ({result.file_path})")
|
|||
|
|
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
data="\n".join(output_parts),
|
|||
|
|
metadata={
|
|||
|
|
"function_name": function_name,
|
|||
|
|
"definition_count": len(context["definition"]),
|
|||
|
|
"callers_count": len(context["callers"]),
|
|||
|
|
"callees_count": len(context["callees"]),
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
return ToolResult(
|
|||
|
|
success=False,
|
|||
|
|
error=f"函数上下文搜索失败: {str(e)}",
|
|||
|
|
)
|
|||
|
|
|