wehub-resource-sync
139 lines
5.5 KiB
Python
139 lines
5.5 KiB
Python
import pixeltable as pxt
|
|
from mcp.server.fastmcp import FastMCP
|
|
from pixeltable.iterators import DocumentSplitter
|
|
from pixeltable.functions.huggingface import sentence_transformer
|
|
|
|
mcp = FastMCP("Pixeltable")
|
|
|
|
# Base directory for all indexes
|
|
DIRECTORY = 'doc_search'
|
|
|
|
# Registry to hold all document indexes
|
|
document_indexes = {}
|
|
|
|
@mcp.tool()
|
|
def setup_document_index(table_name: str) -> str:
|
|
"""Set up a document index with the provided name.
|
|
|
|
Args:
|
|
table_name: The name of the document index (e.g., 'reports', 'articles').
|
|
|
|
Returns:
|
|
A message indicating whether the index was created, already exists, or failed.
|
|
"""
|
|
global document_indexes
|
|
try:
|
|
# Construct full table and view names
|
|
full_table_name = f'{DIRECTORY}.{table_name}'
|
|
chunks_view_name = f'{DIRECTORY}.{table_name}_chunks'
|
|
|
|
# Check if the table already exists
|
|
existing_tables = pxt.list_tables()
|
|
if full_table_name in existing_tables:
|
|
document_index = pxt.get_table(full_table_name)
|
|
chunks_view = pxt.get_table(chunks_view_name)
|
|
document_indexes[full_table_name] = (document_index, chunks_view)
|
|
return f"Document index '{full_table_name}' already exists and is ready for use."
|
|
|
|
# Create directory and table
|
|
pxt.create_dir(DIRECTORY, if_exists='ignore')
|
|
document_index = pxt.create_table(
|
|
full_table_name,
|
|
{'pdf_file': pxt.Document},
|
|
if_exists='ignore'
|
|
)
|
|
|
|
# Create view for document chunks
|
|
chunks_view = pxt.create_view(
|
|
chunks_view_name,
|
|
document_index,
|
|
iterator=DocumentSplitter.create(
|
|
document=document_index.pdf_file,
|
|
separators='token_limit',
|
|
limit=300 # Tokens per chunk
|
|
),
|
|
if_exists='ignore'
|
|
)
|
|
|
|
# Define the embedding model and create embedding index
|
|
embed_model = sentence_transformer.using(model_id='intfloat/e5-large-v2')
|
|
chunks_view.add_embedding_index(
|
|
column='text',
|
|
string_embed=embed_model,
|
|
if_exists='ignore'
|
|
)
|
|
|
|
# Store in the registry
|
|
document_indexes[full_table_name] = (document_index, chunks_view)
|
|
return f"Document index '{full_table_name}' created successfully."
|
|
except Exception as e:
|
|
return f"Error setting up document index '{full_table_name}': {str(e)}"
|
|
|
|
@mcp.tool()
|
|
def insert_document(table_name: str, document_location: str) -> str:
|
|
"""Insert a document file into the specified document index.
|
|
|
|
Args:
|
|
table_name: The name of the document index (e.g., 'reports', 'articles').
|
|
document_location: The URL or path to the document file to insert (e.g., local path or URL).
|
|
|
|
Returns:
|
|
A confirmation message indicating success or failure.
|
|
"""
|
|
full_table_name = f'{DIRECTORY}.{table_name}'
|
|
try:
|
|
if full_table_name not in document_indexes:
|
|
return f"Error: Document index '{full_table_name}' not set up. Please call setup_document_index first."
|
|
document_index, _ = document_indexes[full_table_name]
|
|
document_index.insert([{'pdf_file': document_location}])
|
|
return f"Document file '{document_location}' inserted successfully into index '{full_table_name}'."
|
|
except Exception as e:
|
|
return f"Error inserting document file into '{full_table_name}': {str(e)}"
|
|
|
|
@mcp.tool()
|
|
def query_document(table_name: str, query_text: str, top_n: int = 5) -> str:
|
|
"""Query the specified document index with a text question.
|
|
|
|
Args:
|
|
table_name: The name of the document index (e.g., 'reports', 'articles').
|
|
query_text: The question or text to search for in the document content.
|
|
top_n: Number of top results to return (default is 5).
|
|
|
|
Returns:
|
|
A string containing the top matching text chunks and their similarity scores.
|
|
"""
|
|
full_table_name = f'{DIRECTORY}.{table_name}'
|
|
try:
|
|
if full_table_name not in document_indexes:
|
|
return f"Error: Document index '{full_table_name}' not set up. Please call setup_document_index first."
|
|
_, chunks_view = document_indexes[full_table_name]
|
|
|
|
# Calculate similarity scores
|
|
sim = chunks_view.text.similarity(query_text)
|
|
|
|
# Get top results
|
|
results = (chunks_view.order_by(sim, asc=False)
|
|
.select(chunks_view.text, sim=sim)
|
|
.limit(top_n)
|
|
.collect())
|
|
|
|
# Format the results
|
|
result_str = f"Query Results for '{query_text}' in '{full_table_name}':\n\n"
|
|
for i, row in enumerate(results.to_pandas().itertuples(), 1):
|
|
result_str += f"{i}. Score: {row.sim:.4f}\n"
|
|
result_str += f" Text: {row.text}\n\n"
|
|
|
|
return result_str if result_str else "No results found."
|
|
except Exception as e:
|
|
return f"Error querying document index '{full_table_name}': {str(e)}"
|
|
|
|
@mcp.tool()
|
|
def list_document_tables() -> str:
|
|
"""List all document indexes currently available.
|
|
|
|
Returns:
|
|
A string listing the current document indexes.
|
|
"""
|
|
tables = pxt.list_tables()
|
|
document_tables = [t for t in tables if t.startswith(f'{DIRECTORY}.')]
|
|
return f"Current document indexes: {', '.join(document_tables)}" if document_tables else "No document indexes exist." |