gmakstutis/openrag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add Chat mcp using sdk

Browse source
This commit is contained in:
Edwin Jose 2025-12-26 13:43:06 -05:00
parent bf716b49c8
commit cf8750d906
9 changed files with 1650 additions and 1560 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
pyproject.toml
View file

@ -35,7 +35,8 @@ dependencies = [
"docling-serve==1.5.0",
"docling-core==2.48.1",
"easyocr>=1.7.1; sys_platform != 'darwin'",
"zxcvbn>=4.5.0"
"zxcvbn>=4.5.0",
"openrag-sdk>=0.1.0",
]
[dependency-groups]

1
sdks/mcp/pyproject.toml
View file

@ -19,6 +19,7 @@ classifiers = [
dependencies = [
"mcp>=1.0.0",
"httpx>=0.27.0",
"openrag-sdk>=0.1.0",
]
[project.scripts]

19
sdks/mcp/src/openrag_mcp/config.py
View file

@ -2,6 +2,8 @@
import os
from openrag_sdk import OpenRAGClient
class Config:
"""Configuration loaded from environment variables."""
@ -26,6 +28,7 @@ class Config:
_config: Config | None = None
_openrag_client: OpenRAGClient | None = None
def get_config() -> Config:
@ -36,8 +39,21 @@ def get_config() -> Config:
return _config
def get_openrag_client() -> OpenRAGClient:
"""Get singleton OpenRAGClient instance."""
global _openrag_client
if _openrag_client is None:
# OpenRAGClient reads OPENRAG_API_KEY and OPENRAG_URL from env
_openrag_client = OpenRAGClient()
return _openrag_client
def get_client():
"""Get an httpx async client configured for OpenRAG."""
"""Get an httpx async client configured for OpenRAG.
This is kept for backward compatibility with operations
not yet supported by the SDK (list_documents, ingest_url).
"""
import httpx
config = get_config()
@ -46,4 +62,3 @@ def get_client():
headers=config.headers,
timeout=60.0,
)

8
sdks/mcp/src/openrag_mcp/server.py
View file

@ -8,8 +8,8 @@ from mcp.server.stdio import stdio_server
from openrag_mcp.config import get_config
from openrag_mcp.tools.chat import register_chat_tools
from openrag_mcp.tools.search import register_search_tools
from openrag_mcp.tools.documents import register_document_tools
# from openrag_mcp.tools.search import register_search_tools
# from openrag_mcp.tools.documents import register_document_tools
# Configure logging to stderr (stdout is used for MCP protocol)
logging.basicConfig(
@ -31,8 +31,8 @@ def create_server() -> Server:
# Register all tools
register_chat_tools(server)
register_search_tools(server)
register_document_tools(server)
# register_search_tools(server)
# register_document_tools(server)
logger.info("OpenRAG MCP server initialized with all tools")
return server

6
sdks/mcp/src/openrag_mcp/tools/__init__.py
View file

@ -1,8 +1,8 @@
"""OpenRAG MCP tools."""
from openrag_mcp.tools.chat import register_chat_tools
from openrag_mcp.tools.search import register_search_tools
from openrag_mcp.tools.documents import register_document_tools
# from openrag_mcp.tools.search import register_search_tools
# from openrag_mcp.tools.documents import register_document_tools
__all__ = ["register_chat_tools", "register_search_tools", "register_document_tools"]
__all__ = ["register_chat_tools"]

86
sdks/mcp/src/openrag_mcp/tools/chat.py
View file

@ -1,12 +1,19 @@
"""Chat tool for OpenRAG MCP server."""
import json
import logging
from mcp.server import Server
from mcp.types import TextContent, Tool
from openrag_mcp.config import get_client
from openrag_sdk import (
AuthenticationError,
OpenRAGError,
RateLimitError,
ServerError,
ValidationError,
)
from openrag_mcp.config import get_openrag_client
logger = logging.getLogger("openrag-mcp.chat")
@ -23,7 +30,8 @@ def register_chat_tools(server: Server) -> None:
description=(
"Send a message to OpenRAG and get a RAG-enhanced response. "
"The response is informed by documents in your knowledge base. "
"Use chat_id to continue a previous conversation."
"Use chat_id to continue a previous conversation, or filter_id "
"to apply a knowledge filter."
),
inputSchema={
"type": "object",
@ -36,11 +44,20 @@ def register_chat_tools(server: Server) -> None:
"type": "string",
"description": "Optional conversation ID to continue a previous chat",
},
"filter_id": {
"type": "string",
"description": "Optional knowledge filter ID to apply",
},
"limit": {
"type": "integer",
"description": "Maximum number of sources to retrieve (default: 10)",
"default": 10,
},
"score_threshold": {
"type": "number",
"description": "Minimum relevance score threshold (default: 0)",
"default": 0,
},
},
"required": ["message"],
},
@ -55,46 +72,51 @@ def register_chat_tools(server: Server) -> None:
message = arguments.get("message", "")
chat_id = arguments.get("chat_id")
filter_id = arguments.get("filter_id")
limit = arguments.get("limit", 10)
score_threshold = arguments.get("score_threshold", 0)
if not message:
return [TextContent(type="text", text="Error: message is required")]
try:
async with get_client() as client:
payload = {
"message": message,
"stream": False,
"limit": limit,
}
if chat_id:
payload["chat_id"] = chat_id
client = get_openrag_client()
response = await client.chat.create(
message=message,
chat_id=chat_id,
filter_id=filter_id,
limit=limit,
score_threshold=score_threshold,
)
response = await client.post("/api/v1/chat", json=payload)
response.raise_for_status()
data = response.json()
# Build formatted response
output_parts = [response.response]
# Format the response
result_text = data.get("response", "")
sources = data.get("sources", [])
new_chat_id = data.get("chat_id")
if response.sources:
output_parts.append("\n\n---\n**Sources:**")
for i, source in enumerate(response.sources, 1):
output_parts.append(f"\n{i}. {source.filename} (relevance: {source.score:.2f})")
# Build formatted response
output_parts = [result_text]
if response.chat_id:
output_parts.append(f"\n\n_Chat ID: {response.chat_id}_")
if sources:
output_parts.append("\n\n---\n**Sources:**")
for i, source in enumerate(sources, 1):
filename = source.get("filename", "Unknown")
score = source.get("score", 0)
output_parts.append(f"\n{i}. {filename} (relevance: {score:.2f})")
if new_chat_id:
output_parts.append(f"\n\n_Chat ID: {new_chat_id}_")
return [TextContent(type="text", text="".join(output_parts))]
return [TextContent(type="text", text="".join(output_parts))]
except AuthenticationError as e:
logger.error(f"Authentication error: {e.message}")
return [TextContent(type="text", text=f"Authentication error: {e.message}")]
except ValidationError as e:
logger.error(f"Validation error: {e.message}")
return [TextContent(type="text", text=f"Invalid request: {e.message}")]
except RateLimitError as e:
logger.error(f"Rate limit error: {e.message}")
return [TextContent(type="text", text=f"Rate limited: {e.message}")]
except ServerError as e:
logger.error(f"Server error: {e.message}")
return [TextContent(type="text", text=f"Server error: {e.message}")]
except OpenRAGError as e:
logger.error(f"OpenRAG error: {e.message}")
return [TextContent(type="text", text=f"Error: {e.message}")]
except Exception as e:
logger.error(f"Chat error: {e}")
return [TextContent(type="text", text=f"Error: {str(e)}")]

409
sdks/mcp/src/openrag_mcp/tools/documents.py
View file

@ -1,236 +1,257 @@
"""Document tools for OpenRAG MCP server."""
# """Document tools for OpenRAG MCP server."""
import logging
import os
from pathlib import Path
# import logging
# from pathlib import Path
from mcp.server import Server
from mcp.types import TextContent, Tool
# from mcp.server import Server
# from mcp.types import TextContent, Tool
from openrag_mcp.config import get_client
# from openrag_sdk import (
# AuthenticationError,
# NotFoundError,
# OpenRAGError,
# RateLimitError,
# ServerError,
# ValidationError,
# )
logger = logging.getLogger("openrag-mcp.documents")
# from openrag_mcp.config import get_client, get_openrag_client
# logger = logging.getLogger("openrag-mcp.documents")
def register_document_tools(server: Server) -> None:
"""Register document-related tools with the MCP server."""
# def register_document_tools(server: Server) -> None:
# """Register document-related tools with the MCP server."""
@server.list_tools()
async def list_document_tools() -> list[Tool]:
"""List document tools."""
return [
Tool(
name="openrag_ingest_file",
description=(
"Ingest a local file into the OpenRAG knowledge base. "
"Supported formats: PDF, DOCX, TXT, MD, HTML, and more."
),
inputSchema={
"type": "object",
"properties": {
"file_path": {
"type": "string",
"description": "Absolute path to the file to ingest",
},
},
"required": ["file_path"],
},
),
Tool(
name="openrag_ingest_url",
description=(
"Ingest content from a URL into the OpenRAG knowledge base. "
"The URL content will be fetched, processed, and stored."
),
inputSchema={
"type": "object",
"properties": {
"url": {
"type": "string",
"description": "The URL to fetch and ingest",
},
},
"required": ["url"],
},
),
Tool(
name="openrag_list_documents",
description="List documents in the OpenRAG knowledge base.",
inputSchema={
"type": "object",
"properties": {
"limit": {
"type": "integer",
"description": "Maximum number of documents to return (default: 50)",
"default": 50,
},
},
"required": [],
},
),
Tool(
name="openrag_delete_document",
description="Delete a document from the OpenRAG knowledge base.",
inputSchema={
"type": "object",
"properties": {
"filename": {
"type": "string",
"description": "Name of the file to delete",
},
},
"required": ["filename"],
},
),
]
# @server.list_tools()
# async def list_document_tools() -> list[Tool]:
# """List document tools."""
# return [
# Tool(
# name="openrag_ingest_file",
# description=(
# "Ingest a local file into the OpenRAG knowledge base. "
# "Supported formats: PDF, DOCX, TXT, MD, HTML, and more."
# ),
# inputSchema={
# "type": "object",
# "properties": {
# "file_path": {
# "type": "string",
# "description": "Absolute path to the file to ingest",
# },
# },
# "required": ["file_path"],
# },
# ),
# Tool(
# name="openrag_ingest_url",
# description=(
# "Ingest content from a URL into the OpenRAG knowledge base. "
# "The URL content will be fetched, processed, and stored."
# ),
# inputSchema={
# "type": "object",
# "properties": {
# "url": {
# "type": "string",
# "description": "The URL to fetch and ingest",
# },
# },
# "required": ["url"],
# },
# ),
# Tool(
# name="openrag_list_documents",
# description="List documents in the OpenRAG knowledge base.",
# inputSchema={
# "type": "object",
# "properties": {
# "limit": {
# "type": "integer",
# "description": "Maximum number of documents to return (default: 50)",
# "default": 50,
# },
# },
# "required": [],
# },
# ),
# Tool(
# name="openrag_delete_document",
# description="Delete a document from the OpenRAG knowledge base.",
# inputSchema={
# "type": "object",
# "properties": {
# "filename": {
# "type": "string",
# "description": "Name of the file to delete",
# },
# },
# "required": ["filename"],
# },
# ),
# ]
@server.call_tool()
async def call_document_tool(name: str, arguments: dict) -> list[TextContent]:
"""Handle document tool calls."""
if name == "openrag_ingest_file":
return await _ingest_file(arguments)
elif name == "openrag_ingest_url":
return await _ingest_url(arguments)
elif name == "openrag_list_documents":
return await _list_documents(arguments)
elif name == "openrag_delete_document":
return await _delete_document(arguments)
return []
# @server.call_tool()
# async def call_document_tool(name: str, arguments: dict) -> list[TextContent]:
# """Handle document tool calls."""
# if name == "openrag_ingest_file":
# return await _ingest_file(arguments)
# elif name == "openrag_ingest_url":
# return await _ingest_url(arguments)
# elif name == "openrag_list_documents":
# return await _list_documents(arguments)
# elif name == "openrag_delete_document":
# return await _delete_document(arguments)
# return []
async def _ingest_file(arguments: dict) -> list[TextContent]:
"""Ingest a local file into OpenRAG."""
file_path = arguments.get("file_path", "")
# async def _ingest_file(arguments: dict) -> list[TextContent]:
# """Ingest a local file into OpenRAG using the SDK."""
# file_path = arguments.get("file_path", "")
if not file_path:
return [TextContent(type="text", text="Error: file_path is required")]
# if not file_path:
# return [TextContent(type="text", text="Error: file_path is required")]
path = Path(file_path)
# path = Path(file_path)
if not path.exists():
return [TextContent(type="text", text=f"Error: File not found: {file_path}")]
# if not path.exists():
# return [TextContent(type="text", text=f"Error: File not found: {file_path}")]
if not path.is_file():
return [TextContent(type="text", text=f"Error: Path is not a file: {file_path}")]
# if not path.is_file():
# return [TextContent(type="text", text=f"Error: Path is not a file: {file_path}")]
try:
async with get_client() as client:
# Read file and upload
with open(path, "rb") as f:
files = {"file": (path.name, f)}
# Remove Content-Type header for multipart upload
headers = dict(client.headers)
headers.pop("Content-Type", None)
# try:
# client = get_openrag_client()
# # Use wait=False to return immediately with task_id
# response = await client.documents.ingest(file_path=path, wait=False)
response = await client.post(
"/api/v1/documents/ingest",
files=files,
headers=headers,
)
response.raise_for_status()
data = response.json()
# result = f"Successfully queued '{response.filename or path.name}' for ingestion."
# if response.task_id:
# result += f"\nTask ID: {response.task_id}"
task_id = data.get("task_id")
filename = data.get("filename", path.name)
# return [TextContent(type="text", text=result)]
result = f"Successfully queued '{filename}' for ingestion."
if task_id:
result += f"\nTask ID: {task_id}"
return [TextContent(type="text", text=result)]
except Exception as e:
logger.error(f"Ingest file error: {e}")
return [TextContent(type="text", text=f"Error ingesting file: {str(e)}")]
# except AuthenticationError as e:
# logger.error(f"Authentication error: {e.message}")
# return [TextContent(type="text", text=f"Authentication error: {e.message}")]
# except ValidationError as e:
# logger.error(f"Validation error: {e.message}")
# return [TextContent(type="text", text=f"Invalid request: {e.message}")]
# except RateLimitError as e:
# logger.error(f"Rate limit error: {e.message}")
# return [TextContent(type="text", text=f"Rate limited: {e.message}")]
# except ServerError as e:
# logger.error(f"Server error: {e.message}")
# return [TextContent(type="text", text=f"Server error: {e.message}")]
# except OpenRAGError as e:
# logger.error(f"OpenRAG error: {e.message}")
# return [TextContent(type="text", text=f"Error: {e.message}")]
# except Exception as e:
# logger.error(f"Ingest file error: {e}")
# return [TextContent(type="text", text=f"Error ingesting file: {str(e)}")]
async def _ingest_url(arguments: dict) -> list[TextContent]:
"""Ingest content from a URL into OpenRAG."""
url = arguments.get("url", "")
# async def _ingest_url(arguments: dict) -> list[TextContent]:
# """Ingest content from a URL into OpenRAG.
if not url:
return [TextContent(type="text", text="Error: url is required")]
# Note: This uses the SDK's chat to trigger URL ingestion via the agent.
# """
# url = arguments.get("url", "")
if not url.startswith(("http://", "https://")):
return [TextContent(type="text", text="Error: url must start with http:// or https://")]
# if not url:
# return [TextContent(type="text", text="Error: url is required")]
try:
# Use chat with a special prompt to trigger URL ingestion via the agent
async with get_client() as client:
payload = {
"message": f"Please ingest the content from this URL into the knowledge base: {url}",
"stream": False,
}
# if not url.startswith(("http://", "https://")):
# return [TextContent(type="text", text="Error: url must start with http:// or https://")]
response = await client.post("/api/v1/chat", json=payload)
response.raise_for_status()
data = response.json()
# try:
# # Use chat with a special prompt to trigger URL ingestion via the agent
# client = get_openrag_client()
# response = await client.chat.create(
# message=f"Please ingest the content from this URL into the knowledge base: {url}",
# )
result_text = data.get("response", "")
return [TextContent(type="text", text=f"URL ingestion requested.\n\n{result_text}")]
# return [TextContent(type="text", text=f"URL ingestion requested.\n\n{response.response}")]
except Exception as e:
logger.error(f"Ingest URL error: {e}")
return [TextContent(type="text", text=f"Error ingesting URL: {str(e)}")]
# except AuthenticationError as e:
# logger.error(f"Authentication error: {e.message}")
# return [TextContent(type="text", text=f"Authentication error: {e.message}")]
# except ServerError as e:
# logger.error(f"Server error: {e.message}")
# return [TextContent(type="text", text=f"Server error: {e.message}")]
# except OpenRAGError as e:
# logger.error(f"OpenRAG error: {e.message}")
# return [TextContent(type="text", text=f"Error: {e.message}")]
# except Exception as e:
# logger.error(f"Ingest URL error: {e}")
# return [TextContent(type="text", text=f"Error ingesting URL: {str(e)}")]
async def _list_documents(arguments: dict) -> list[TextContent]:
"""List documents in the knowledge base."""
limit = arguments.get("limit", 50)
# async def _list_documents(arguments: dict) -> list[TextContent]:
# """List documents in the knowledge base.
try:
async with get_client() as client:
response = await client.get("/api/v1/documents", params={"limit": limit})
response.raise_for_status()
data = response.json()
# Note: This uses direct HTTP calls as the SDK doesn't yet support listing documents.
# """
# limit = arguments.get("limit", 50)
documents = data.get("documents", [])
# try:
# async with get_client() as client:
# response = await client.get("/api/v1/documents", params={"limit": limit})
# response.raise_for_status()
# data = response.json()
if not documents:
return [TextContent(type="text", text="No documents found in the knowledge base.")]
# documents = data.get("documents", [])
output_parts = [f"Found {len(documents)} document(s):\n"]
# if not documents:
# return [TextContent(type="text", text="No documents found in the knowledge base.")]
for doc in documents:
filename = doc.get("filename", "Unknown")
chunks = doc.get("chunk_count", 0)
created = doc.get("created_at", "")
# output_parts = [f"Found {len(documents)} document(s):\n"]
output_parts.append(f"\n- **{filename}** ({chunks} chunks)")
if created:
output_parts.append(f" - Added: {created[:10]}")
# for doc in documents:
# filename = doc.get("filename", "Unknown")
# chunks = doc.get("chunk_count", 0)
# created = doc.get("created_at", "")
return [TextContent(type="text", text="".join(output_parts))]
# output_parts.append(f"\n- **{filename}** ({chunks} chunks)")
# if created:
# output_parts.append(f" - Added: {created[:10]}")
except Exception as e:
logger.error(f"List documents error: {e}")
return [TextContent(type="text", text=f"Error listing documents: {str(e)}")]
# return [TextContent(type="text", text="".join(output_parts))]
# except Exception as e:
# logger.error(f"List documents error: {e}")
# return [TextContent(type="text", text=f"Error listing documents: {str(e)}")]
async def _delete_document(arguments: dict) -> list[TextContent]:
"""Delete a document from the knowledge base."""
filename = arguments.get("filename", "")
# async def _delete_document(arguments: dict) -> list[TextContent]:
# """Delete a document from the knowledge base using the SDK."""
# filename = arguments.get("filename", "")
if not filename:
return [TextContent(type="text", text="Error: filename is required")]
# if not filename:
# return [TextContent(type="text", text="Error: filename is required")]
try:
async with get_client() as client:
response = await client.request(
"DELETE",
"/api/v1/documents",
json={"filename": filename},
)
response.raise_for_status()
data = response.json()
# try:
# client = get_openrag_client()
# response = await client.documents.delete(filename)
deleted_count = data.get("deleted_count", 0)
return [TextContent(
type="text",
text=f"Successfully deleted '{filename}' ({deleted_count} chunks removed).",
)]
except Exception as e:
logger.error(f"Delete document error: {e}")
return [TextContent(type="text", text=f"Error deleting document: {str(e)}")]
# return [TextContent(
# type="text",
# text=f"Successfully deleted '{filename}' ({response.deleted_chunks} chunks removed).",
# )]
# except NotFoundError as e:
# logger.error(f"Document not found: {e.message}")
# return [TextContent(type="text", text=f"Document not found: {e.message}")]
# except AuthenticationError as e:
# logger.error(f"Authentication error: {e.message}")
# return [TextContent(type="text", text=f"Authentication error: {e.message}")]
# except ServerError as e:
# logger.error(f"Server error: {e.message}")
# return [TextContent(type="text", text=f"Server error: {e.message}")]
# except OpenRAGError as e:
# logger.error(f"OpenRAG error: {e.message}")
# return [TextContent(type="text", text=f"Error: {e.message}")]
# except Exception as e:
# logger.error(f"Delete document error: {e}")
# return [TextContent(type="text", text=f"Error deleting document: {str(e)}")]

15
sdks/mcp/uv.lock generated
View file

@ -337,12 +337,27 @@ source = { editable = "." }
dependencies = [
{ name = "httpx" },
{ name = "mcp" },
{ name = "openrag-sdk" },
]
[package.metadata]
requires-dist = [
{ name = "httpx", specifier = ">=0.27.0" },
{ name = "mcp", specifier = ">=1.0.0" },
{ name = "openrag-sdk", specifier = ">=0.1.0" },
]
[[package]]
name = "openrag-sdk"
version = "0.1.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "httpx" },
{ name = "pydantic" },
]
sdist = { url = "https://files.pythonhosted.org/packages/d7/9e/7a10ddb6742417970163d11aed40146bb3340532f17c1099434d8667a4e0/openrag_sdk-0.1.0.tar.gz", hash = "sha256:cf99fddf254c6c72295c73498877c56f68287d33a072ca2171f490de7df8617e", size = 17116 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/74/3c/6962d87b8b4604ef8cf776bf7d4e416a0025d5f08f5400ecc47b35616ae0/openrag_sdk-0.1.0-py3-none-any.whl", hash = "sha256:6b99be26b64a61e2347a4ce8f530b673a119d59bfe11ff5ead6318b9a3fb9dac", size = 15756 },
]
[[package]]

2663
uv.lock generated
View file

File diff suppressed because it is too large Load diff