"""Redis 벡터 DB MCP 도구들"""
from typing import Any, List
from mcp import types
from ..utils.redis_vector import save_document, get_document, qa_with_redis
class SaveDocumentTool:
"""문서를 벡터 DB에 저장하는 MCP 도구"""
@staticmethod
def get_tool_definition() -> types.Tool:
"""도구 정의 반환"""
return types.Tool(
name="save_document",
description="공식 문서를 벡터 DB에 저장하고 생성된 vector_id를 반환합니다.",
inputSchema={
"type": "object",
"properties": {
"title": {
"type": "string",
"description": "문서 제목"
},
"url": {
"type": "string",
"description": "문서 URL"
},
"content": {
"type": "string",
"description": "문서 내용"
}
},
"required": ["title", "url", "content"]
}
)
@staticmethod
async def execute(arguments: dict[str, Any]) -> List[types.TextContent]:
"""도구 실행"""
title = arguments.get("title", "")
url = arguments.get("url", "")
content = arguments.get("content", "")
if not all([title, url, content]):
return [types.TextContent(
type="text",
text="title, url, content를 모두 입력해주세요."
)]
try:
vector_id = save_document(title, url, content)
return [types.TextContent(
type="text",
text=f"문서가 저장되었습니다.\nvector_id: {vector_id}"
)]
except Exception as e:
return [types.TextContent(
type="text",
text=f"오류가 발생했습니다: {str(e)}"
)]
class GetDocumentTool:
"""vector_id로 문서를 조회하는 MCP 도구"""
@staticmethod
def get_tool_definition() -> types.Tool:
"""도구 정의 반환"""
return types.Tool(
name="get_document",
description="vector_id로 저장된 문서를 조회합니다.",
inputSchema={
"type": "object",
"properties": {
"vector_id": {
"type": "string",
"description": "조회할 문서의 vector_id"
}
},
"required": ["vector_id"]
}
)
@staticmethod
async def execute(arguments: dict[str, Any]) -> List[types.TextContent]:
"""도구 실행"""
vector_id = arguments.get("vector_id", "")
if not vector_id:
return [types.TextContent(
type="text",
text="vector_id를 입력해주세요."
)]
try:
doc = get_document(vector_id)
if not doc:
return [types.TextContent(
type="text",
text=f"vector_id '{vector_id}'에 해당하는 문서를 찾을 수 없습니다."
)]
# bytes를 문자열로 변환
title = doc[b"title"].decode() if isinstance(doc[b"title"], bytes) else doc[b"title"]
url = doc[b"url"].decode() if isinstance(doc[b"url"], bytes) else doc[b"url"]
content = doc[b"content"].decode() if isinstance(doc[b"content"], bytes) else doc[b"content"]
result = f"""문서 정보:
제목: {title}
URL: {url}
내용: {content[:200]}{'...' if len(content) > 200 else ''}"""
return [types.TextContent(
type="text",
text=result
)]
except Exception as e:
return [types.TextContent(
type="text",
text=f"오류가 발생했습니다: {str(e)}"
)]
class QAWithRedisTool:
"""vector_id에 저장된 문서를 기반으로 질문에 답변하는 MCP 도구"""
@staticmethod
def get_tool_definition() -> types.Tool:
"""도구 정의 반환"""
return types.Tool(
name="qa_with_redis",
description="vector_id에 저장된 문서를 기반으로 질문에 대한 답변을 생성합니다.",
inputSchema={
"type": "object",
"properties": {
"vector_id": {
"type": "string",
"description": "참조할 문서의 vector_id"
},
"query": {
"type": "string",
"description": "질문"
}
},
"required": ["vector_id", "query"]
}
)
@staticmethod
async def execute(arguments: dict[str, Any]) -> List[types.TextContent]:
"""도구 실행"""
vector_id = arguments.get("vector_id", "")
query = arguments.get("query", "")
if not all([vector_id, query]):
return [types.TextContent(
type="text",
text="vector_id와 query를 모두 입력해주세요."
)]
try:
answer = qa_with_redis(vector_id, query)
return [types.TextContent(
type="text",
text=answer
)]
except Exception as e:
return [types.TextContent(
type="text",
text=f"오류가 발생했습니다: {str(e)}"
)]
