"""MCP tools for searching statements."""
from ..database.lancedb_client import LanceDBClient
from ..embeddings.generator import EmbeddingGenerator
import json
class SearchTools:
"""Tools for semantic search over statements."""
def __init__(self, lance_client: LanceDBClient, embedding_gen: EmbeddingGenerator):
"""Initialize search tools.
Args:
lance_client: LanceDB client
embedding_gen: Embedding generator
"""
self.lance = lance_client
self.embedder = embedding_gen
async def search_statements(self, query: str, limit: int = 10) -> list[dict]:
"""Search statements using natural language.
Args:
query: Natural language query
limit: Maximum number of results
Returns:
List of matching results
"""
# Generate query embedding
query_embedding = self.embedder.encode(query)
# Search in LanceDB
results = await self.lance.search(
query_vector=query_embedding,
limit=limit,
)
# Format results
formatted_results = []
for result in results:
formatted_results.append({
"statement_id": result.get("statement_id"),
"content": result.get("content"),
"chunk_type": result.get("chunk_type"),
"account_number": result.get("account_number"),
"statement_date": result.get("statement_date"),
"institution": result.get("institution"),
"score": result.get("_distance", 0.0), # Lower is better
})
return formatted_results
