rag_tools.py•2.49 kB
import os
from typing import Any
from rag_service import RagService
class RagTools:
"""Tools class that defines MCP tools for RAG operations, equivalent to the Java RagTools"""
def __init__(self):
self.user_id_from_environment = os.getenv('USER_ID', 'invalid')
self.rag_service = RagService()
def list_content_names(self, name: str | None = None) -> list[dict[str, Any]]:
"""
List all content names from the organization's database optionally filtered by name.
Args:
name: Optional name filter to search for specific content
Returns:
List of content objects
"""
return self.rag_service.get_all_contents(name, self.user_id_from_environment)
def search_organization_contents(self, query: str) -> list[str]:
"""
Search through the organization's content database using semantic search.
Args:
query: The search query (required)
Returns:
List of search results
"""
return self.rag_service.search_contents(query, self.user_id_from_environment)
def delete_organization_content(self, content_id: str) -> list[str]:
"""
Delete specific content from the organization's database.
Args:
content_id: The ID of the content to delete (required)
Returns:
Success message
"""
self.rag_service.delete_content(content_id, self.user_id_from_environment)
return ["Content deleted successfully"]
def upload_content_file_about_organization(self, file: str, file_name: str, role: str = "") -> list[str]:
"""
Upload content file about the organization.
Args:
file: The file content to upload
file_name: The file name (required)
role: The roles of the user
Returns:
Success message
"""
self.rag_service.upload_file(file, file_name, role, self.user_id_from_environment)
return ["File uploaded successfully"]
def upload_content_url_about_organization(self, url: str, role: str = "") -> list[str]:
"""
Upload content url about the organization.
Args:
url: The URL to upload
role: The roles of the user
Returns:
Success message
"""
self.rag_service.upload_url(url, role, self.user_id_from_environment)
return ["URL uploaded successfully"]