Split out main module
This commit is contained in:
13
src/knowledge_search_mcp/services/__init__.py
Normal file
13
src/knowledge_search_mcp/services/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
"""Service modules for business logic."""
|
||||
|
||||
from .search import filter_search_results, format_search_results, generate_query_embedding
|
||||
from .validation import validate_genai_access, validate_gcs_access, validate_vector_search_access
|
||||
|
||||
__all__ = [
|
||||
"filter_search_results",
|
||||
"format_search_results",
|
||||
"generate_query_embedding",
|
||||
"validate_genai_access",
|
||||
"validate_gcs_access",
|
||||
"validate_vector_search_access",
|
||||
]
|
||||
110
src/knowledge_search_mcp/services/search.py
Normal file
110
src/knowledge_search_mcp/services/search.py
Normal file
@@ -0,0 +1,110 @@
|
||||
# ruff: noqa: INP001
|
||||
"""Search helper functions."""
|
||||
|
||||
from google import genai
|
||||
from google.genai import types as genai_types
|
||||
|
||||
from ..logging import log_structured_entry
|
||||
from ..models import SearchResult
|
||||
|
||||
|
||||
async def generate_query_embedding(
|
||||
genai_client: genai.Client,
|
||||
embedding_model: str,
|
||||
query: str,
|
||||
) -> tuple[list[float], str | None]:
|
||||
"""Generate embedding for search query.
|
||||
|
||||
Returns:
|
||||
Tuple of (embedding vector, error message). Error message is None on success.
|
||||
"""
|
||||
if not query or not query.strip():
|
||||
return ([], "Error: Query cannot be empty")
|
||||
|
||||
log_structured_entry("Generating query embedding", "INFO")
|
||||
try:
|
||||
response = await genai_client.aio.models.embed_content(
|
||||
model=embedding_model,
|
||||
contents=query,
|
||||
config=genai_types.EmbedContentConfig(
|
||||
task_type="RETRIEVAL_QUERY",
|
||||
),
|
||||
)
|
||||
embedding = response.embeddings[0].values
|
||||
return (embedding, None)
|
||||
except Exception as e:
|
||||
error_type = type(e).__name__
|
||||
error_msg = str(e)
|
||||
|
||||
# Check if it's a rate limit error
|
||||
if "429" in error_msg or "RESOURCE_EXHAUSTED" in error_msg:
|
||||
log_structured_entry(
|
||||
"Rate limit exceeded while generating embedding",
|
||||
"WARNING",
|
||||
{
|
||||
"error": error_msg,
|
||||
"error_type": error_type,
|
||||
"query": query[:100]
|
||||
}
|
||||
)
|
||||
return ([], "Error: API rate limit exceeded. Please try again later.")
|
||||
else:
|
||||
log_structured_entry(
|
||||
"Failed to generate query embedding",
|
||||
"ERROR",
|
||||
{
|
||||
"error": error_msg,
|
||||
"error_type": error_type,
|
||||
"query": query[:100]
|
||||
}
|
||||
)
|
||||
return ([], f"Error generating embedding: {error_msg}")
|
||||
|
||||
|
||||
def filter_search_results(
|
||||
results: list[SearchResult],
|
||||
min_similarity: float = 0.6,
|
||||
top_percent: float = 0.9,
|
||||
) -> list[SearchResult]:
|
||||
"""Filter search results by similarity thresholds.
|
||||
|
||||
Args:
|
||||
results: Raw search results from vector search.
|
||||
min_similarity: Minimum similarity score (distance) to include.
|
||||
top_percent: Keep results within this percentage of the top score.
|
||||
|
||||
Returns:
|
||||
Filtered list of search results.
|
||||
"""
|
||||
if not results:
|
||||
return []
|
||||
|
||||
max_sim = max(r["distance"] for r in results)
|
||||
cutoff = max_sim * top_percent
|
||||
|
||||
filtered = [
|
||||
s
|
||||
for s in results
|
||||
if s["distance"] > cutoff and s["distance"] > min_similarity
|
||||
]
|
||||
|
||||
return filtered
|
||||
|
||||
|
||||
def format_search_results(results: list[SearchResult]) -> str:
|
||||
"""Format search results as XML-like documents.
|
||||
|
||||
Args:
|
||||
results: List of search results to format.
|
||||
|
||||
Returns:
|
||||
Formatted string with document tags.
|
||||
"""
|
||||
if not results:
|
||||
return "No relevant documents found for your query."
|
||||
|
||||
formatted_results = [
|
||||
f"<document {i} name={result['id']}>\n{result['content']}\n</document {i}>"
|
||||
for i, result in enumerate(results, start=1)
|
||||
]
|
||||
return "\n".join(formatted_results)
|
||||
171
src/knowledge_search_mcp/services/validation.py
Normal file
171
src/knowledge_search_mcp/services/validation.py
Normal file
@@ -0,0 +1,171 @@
|
||||
# ruff: noqa: INP001
|
||||
"""Validation functions for Google Cloud services."""
|
||||
|
||||
from gcloud.aio.auth import Token
|
||||
from google import genai
|
||||
from google.genai import types as genai_types
|
||||
|
||||
from ..clients.vector_search import GoogleCloudVectorSearch
|
||||
from ..config import Settings
|
||||
from ..logging import log_structured_entry
|
||||
|
||||
|
||||
async def validate_genai_access(genai_client: genai.Client, cfg: Settings) -> str | None:
|
||||
"""Validate GenAI embedding access.
|
||||
|
||||
Returns:
|
||||
Error message if validation fails, None if successful.
|
||||
"""
|
||||
log_structured_entry("Validating GenAI embedding access", "INFO")
|
||||
try:
|
||||
test_response = await genai_client.aio.models.embed_content(
|
||||
model=cfg.embedding_model,
|
||||
contents="test",
|
||||
config=genai_types.EmbedContentConfig(
|
||||
task_type="RETRIEVAL_QUERY",
|
||||
),
|
||||
)
|
||||
if test_response and test_response.embeddings:
|
||||
embedding_values = test_response.embeddings[0].values
|
||||
log_structured_entry(
|
||||
"GenAI embedding validation successful",
|
||||
"INFO",
|
||||
{"embedding_dimension": len(embedding_values) if embedding_values else 0}
|
||||
)
|
||||
return None
|
||||
else:
|
||||
msg = "Embedding validation returned empty response"
|
||||
log_structured_entry(msg, "WARNING")
|
||||
return msg
|
||||
except Exception as e:
|
||||
log_structured_entry(
|
||||
"Failed to validate GenAI embedding access - service may not work correctly",
|
||||
"WARNING",
|
||||
{"error": str(e), "error_type": type(e).__name__}
|
||||
)
|
||||
return f"GenAI: {str(e)}"
|
||||
|
||||
|
||||
async def validate_gcs_access(vs: GoogleCloudVectorSearch, cfg: Settings) -> str | None:
|
||||
"""Validate GCS bucket access.
|
||||
|
||||
Returns:
|
||||
Error message if validation fails, None if successful.
|
||||
"""
|
||||
log_structured_entry(
|
||||
"Validating GCS bucket access",
|
||||
"INFO",
|
||||
{"bucket": cfg.bucket}
|
||||
)
|
||||
try:
|
||||
session = vs.storage._get_aio_session()
|
||||
token_obj = Token(
|
||||
session=session,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
access_token = await token_obj.get()
|
||||
headers = {"Authorization": f"Bearer {access_token}"}
|
||||
|
||||
async with session.get(
|
||||
f"https://storage.googleapis.com/storage/v1/b/{cfg.bucket}/o?maxResults=1",
|
||||
headers=headers,
|
||||
) as response:
|
||||
if response.status == 403:
|
||||
msg = f"Access denied to bucket '{cfg.bucket}'. Check permissions."
|
||||
log_structured_entry(
|
||||
"GCS bucket validation failed - access denied - service may not work correctly",
|
||||
"WARNING",
|
||||
{"bucket": cfg.bucket, "status": response.status}
|
||||
)
|
||||
return msg
|
||||
elif response.status == 404:
|
||||
msg = f"Bucket '{cfg.bucket}' not found. Check bucket name and project."
|
||||
log_structured_entry(
|
||||
"GCS bucket validation failed - not found - service may not work correctly",
|
||||
"WARNING",
|
||||
{"bucket": cfg.bucket, "status": response.status}
|
||||
)
|
||||
return msg
|
||||
elif not response.ok:
|
||||
body = await response.text()
|
||||
msg = f"Failed to access bucket '{cfg.bucket}': {response.status}"
|
||||
log_structured_entry(
|
||||
"GCS bucket validation failed - service may not work correctly",
|
||||
"WARNING",
|
||||
{"bucket": cfg.bucket, "status": response.status, "response": body}
|
||||
)
|
||||
return msg
|
||||
else:
|
||||
log_structured_entry(
|
||||
"GCS bucket validation successful",
|
||||
"INFO",
|
||||
{"bucket": cfg.bucket}
|
||||
)
|
||||
return None
|
||||
except Exception as e:
|
||||
log_structured_entry(
|
||||
"Failed to validate GCS bucket access - service may not work correctly",
|
||||
"WARNING",
|
||||
{"error": str(e), "error_type": type(e).__name__, "bucket": cfg.bucket}
|
||||
)
|
||||
return f"GCS: {str(e)}"
|
||||
|
||||
|
||||
async def validate_vector_search_access(vs: GoogleCloudVectorSearch, cfg: Settings) -> str | None:
|
||||
"""Validate vector search endpoint access.
|
||||
|
||||
Returns:
|
||||
Error message if validation fails, None if successful.
|
||||
"""
|
||||
log_structured_entry(
|
||||
"Validating vector search endpoint access",
|
||||
"INFO",
|
||||
{"endpoint_name": cfg.endpoint_name}
|
||||
)
|
||||
try:
|
||||
headers = await vs._async_get_auth_headers()
|
||||
session = vs._get_aio_session()
|
||||
endpoint_url = (
|
||||
f"https://{cfg.location}-aiplatform.googleapis.com/v1/{cfg.endpoint_name}"
|
||||
)
|
||||
|
||||
async with session.get(endpoint_url, headers=headers) as response:
|
||||
if response.status == 403:
|
||||
msg = f"Access denied to endpoint '{cfg.endpoint_name}'. Check permissions."
|
||||
log_structured_entry(
|
||||
"Vector search endpoint validation failed - access denied - service may not work correctly",
|
||||
"WARNING",
|
||||
{"endpoint": cfg.endpoint_name, "status": response.status}
|
||||
)
|
||||
return msg
|
||||
elif response.status == 404:
|
||||
msg = f"Endpoint '{cfg.endpoint_name}' not found. Check endpoint name and project."
|
||||
log_structured_entry(
|
||||
"Vector search endpoint validation failed - not found - service may not work correctly",
|
||||
"WARNING",
|
||||
{"endpoint": cfg.endpoint_name, "status": response.status}
|
||||
)
|
||||
return msg
|
||||
elif not response.ok:
|
||||
body = await response.text()
|
||||
msg = f"Failed to access endpoint '{cfg.endpoint_name}': {response.status}"
|
||||
log_structured_entry(
|
||||
"Vector search endpoint validation failed - service may not work correctly",
|
||||
"WARNING",
|
||||
{"endpoint": cfg.endpoint_name, "status": response.status, "response": body}
|
||||
)
|
||||
return msg
|
||||
else:
|
||||
log_structured_entry(
|
||||
"Vector search endpoint validation successful",
|
||||
"INFO",
|
||||
{"endpoint": cfg.endpoint_name}
|
||||
)
|
||||
return None
|
||||
except Exception as e:
|
||||
log_structured_entry(
|
||||
"Failed to validate vector search endpoint access - service may not work correctly",
|
||||
"WARNING",
|
||||
{"error": str(e), "error_type": type(e).__name__, "endpoint": cfg.endpoint_name}
|
||||
)
|
||||
return f"Vector Search: {str(e)}"
|
||||
Reference in New Issue
Block a user