RAGFlow MCP

import json from configs.ragflow import ragflow from configs.logger import get_logger logger = get_logger(__name__) def get_dataset_by_name(user_id: str): try: datasets = ragflow.list_datasets(name=user_id) return datasets except Exception as e: logger.error(f"Error getting dataset by name {user_id}: {e}") return [] def create_initial_dataset(user_id: str): existed_datas = get_dataset_by_name(user_id) if len(existed_datas) > 0: return existed_datas[0] dataset = ragflow.create_dataset(name=user_id) documents = dataset.upload_documents([{ "display_name": "base_knowledge.txt", "blob": "You are a helpful assistant that can answer questions about the user's data." }]) document_ids = [doc.id for doc in documents] dataset.async_parse_documents(document_ids) return dataset def upload_documents_to_dataset(dataset_name: str, display_names: list[str], blobs: list[str]) -> str: """Uploads documents to a RAGFlow dataset. Args: dataset_name (str): The name of the dataset to upload documents to display_names (list[str]): List of display names for the documents blobs (list[str]): List of document contents as strings Returns: str: Response from the API indicating success or failure """ try: # Get the dataset object using dataset name dataset = ragflow.get_dataset(name=dataset_name) # Prepare documents list documents = [] for display_name, blob in zip(display_names, blobs): documents.append({ "display_name": display_name, "blob": blob }) # Upload documents response = dataset.upload_documents(documents) # Get document IDs doc_info = [] for doc in response: doc_info.append({ "name": doc.display_name if hasattr(doc, 'display_name') else display_names[0], "id": doc.id }) return { "status": "success", "message": f"Successfully uploaded {len(documents)} documents", "dataset": dataset_name, "documents": doc_info } except Exception as e: return { "status": "error", "message": str(e) } def list_datasets( page: int = 1, page_size: int = 30, orderby: str = "create_time", desc: bool = True, name: str = None, dataset_id: str = None ) -> str: """Lists datasets with filtering options. Args: page (int, optional): Page number for pagination. Defaults to 1. page_size (int, optional): Number of items per page. Defaults to 30. orderby (str, optional): Field to sort by ('create_time' or 'update_time'). Defaults to 'create_time'. desc (bool, optional): Sort in descending order if True. Defaults to True. name (str, optional): Filter by dataset name. Defaults to None. dataset_id (str, optional): Filter by dataset ID. Defaults to None. Returns: str: List of datasets matching the criteria """ try: # Validate orderby parameter if orderby not in ["create_time", "update_time"]: return "Error: orderby must be either 'create_time' or 'update_time'" # Build query parameters params = { "page": page, "page_size": page_size, "orderby": orderby, "desc": desc } # Add optional filters if provided if name: params["name"] = name if dataset_id: params["id"] = dataset_id # Get datasets with filters datasets = ragflow.list_datasets(**params) return f"Successfully retrieved datasets: {datasets}" except Exception as e: return f"Failed to list datasets: {str(e)}" def list_documents_in_dataset(dataset_name: str, keywords: str = None) -> str: """Lists documents in a RAGFlow dataset with optional keyword filtering. Args: dataset_name (str): The name of the dataset to list documents from keywords (str, optional): Keywords to filter documents. Defaults to None. Returns: str: List of documents matching the criteria """ try: # Get the dataset object dataset = ragflow.get_dataset(name=dataset_name) # List documents with optional keyword filter documents = dataset.list_documents(keywords=keywords) # Format document information doc_list = [] for doc in documents: doc_list.append({ "id": doc.id, "name": doc.display_name if hasattr(doc, 'display_name') else "Unknown", "status": "listed" }) return { "status": "success", "message": f"Found {len(doc_list)} documents", "dataset": dataset_name, "documents": doc_list } except Exception as e: return { "status": "error", "message": str(e) } def parse_documents_in_dataset(dataset_name: str, document_ids: list[str]) -> str: """Initiates async parsing of documents in a RAGFlow dataset. Args: dataset_name (str): The name of the dataset containing the documents document_ids (list[str]): List of document IDs to parse Returns: str: Response indicating parsing initiation status """ try: # Get the dataset object dataset = ragflow.get_dataset(name=dataset_name) # Initiate async parsing dataset.async_parse_documents(document_ids) return { "status": "success", "message": "Async document parsing initiated", "dataset": dataset_name, "document_count": len(document_ids), "document_ids": document_ids } except Exception as e: return { "status": "error", "message": str(e) }