Add CI
This commit is contained in:
@@ -1,13 +1,21 @@
|
||||
"""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
|
||||
from .search import (
|
||||
filter_search_results,
|
||||
format_search_results,
|
||||
generate_query_embedding,
|
||||
)
|
||||
from .validation import (
|
||||
validate_gcs_access,
|
||||
validate_genai_access,
|
||||
validate_vector_search_access,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"filter_search_results",
|
||||
"format_search_results",
|
||||
"generate_query_embedding",
|
||||
"validate_genai_access",
|
||||
"validate_gcs_access",
|
||||
"validate_genai_access",
|
||||
"validate_vector_search_access",
|
||||
]
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
# 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
|
||||
from knowledge_search_mcp.logging import log_structured_entry
|
||||
from knowledge_search_mcp.models import SearchResult
|
||||
|
||||
|
||||
async def generate_query_embedding(
|
||||
@@ -17,6 +16,7 @@ async def generate_query_embedding(
|
||||
|
||||
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")
|
||||
@@ -30,9 +30,11 @@ async def generate_query_embedding(
|
||||
task_type="RETRIEVAL_QUERY",
|
||||
),
|
||||
)
|
||||
if not response.embeddings or not response.embeddings[0].values:
|
||||
return ([], "Error: Failed to generate embedding - empty response")
|
||||
embedding = response.embeddings[0].values
|
||||
return (embedding, None)
|
||||
except Exception as e:
|
||||
return (embedding, None) # noqa: TRY300
|
||||
except Exception as e: # noqa: BLE001
|
||||
error_type = type(e).__name__
|
||||
error_msg = str(e)
|
||||
|
||||
@@ -41,24 +43,15 @@ async def generate_query_embedding(
|
||||
log_structured_entry(
|
||||
"Rate limit exceeded while generating embedding",
|
||||
"WARNING",
|
||||
{
|
||||
"error": error_msg,
|
||||
"error_type": error_type,
|
||||
"query": query[:100]
|
||||
}
|
||||
{"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}")
|
||||
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(
|
||||
@@ -75,6 +68,7 @@ def filter_search_results(
|
||||
|
||||
Returns:
|
||||
Filtered list of search results.
|
||||
|
||||
"""
|
||||
if not results:
|
||||
return []
|
||||
@@ -82,14 +76,10 @@ def filter_search_results(
|
||||
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 [
|
||||
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.
|
||||
@@ -99,6 +89,7 @@ def format_search_results(results: list[SearchResult]) -> str:
|
||||
|
||||
Returns:
|
||||
Formatted string with document tags.
|
||||
|
||||
"""
|
||||
if not results:
|
||||
return "No relevant documents found for your query."
|
||||
|
||||
@@ -1,20 +1,26 @@
|
||||
# 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
|
||||
from knowledge_search_mcp.clients.vector_search import GoogleCloudVectorSearch
|
||||
from knowledge_search_mcp.config import Settings
|
||||
from knowledge_search_mcp.logging import log_structured_entry
|
||||
|
||||
# HTTP status codes
|
||||
HTTP_FORBIDDEN = 403
|
||||
HTTP_NOT_FOUND = 404
|
||||
|
||||
|
||||
async def validate_genai_access(genai_client: genai.Client, cfg: Settings) -> str | None:
|
||||
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:
|
||||
@@ -30,20 +36,26 @@ async def validate_genai_access(genai_client: genai.Client, cfg: Settings) -> st
|
||||
log_structured_entry(
|
||||
"GenAI embedding validation successful",
|
||||
"INFO",
|
||||
{"embedding_dimension": len(embedding_values) if embedding_values else 0}
|
||||
{
|
||||
"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:
|
||||
msg = "Embedding validation returned empty response"
|
||||
log_structured_entry(msg, "WARNING")
|
||||
return msg # noqa: TRY300
|
||||
except Exception as e: # noqa: BLE001
|
||||
log_structured_entry(
|
||||
"Failed to validate GenAI embedding access - service may not work correctly",
|
||||
(
|
||||
"Failed to validate GenAI embedding access - "
|
||||
"service may not work correctly"
|
||||
),
|
||||
"WARNING",
|
||||
{"error": str(e), "error_type": type(e).__name__}
|
||||
{"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
return f"GenAI: {str(e)}"
|
||||
return f"GenAI: {e!s}"
|
||||
|
||||
|
||||
async def validate_gcs_access(vs: GoogleCloudVectorSearch, cfg: Settings) -> str | None:
|
||||
@@ -51,14 +63,11 @@ async def validate_gcs_access(vs: GoogleCloudVectorSearch, cfg: Settings) -> str
|
||||
|
||||
Returns:
|
||||
Error message if validation fails, None if successful.
|
||||
|
||||
"""
|
||||
log_structured_entry(
|
||||
"Validating GCS bucket access",
|
||||
"INFO",
|
||||
{"bucket": cfg.bucket}
|
||||
)
|
||||
log_structured_entry("Validating GCS bucket access", "INFO", {"bucket": cfg.bucket})
|
||||
try:
|
||||
session = vs.storage._get_aio_session()
|
||||
session = vs.storage._get_aio_session() # noqa: SLF001
|
||||
token_obj = Token(
|
||||
session=session,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
@@ -70,102 +79,136 @@ async def validate_gcs_access(vs: GoogleCloudVectorSearch, cfg: Settings) -> str
|
||||
f"https://storage.googleapis.com/storage/v1/b/{cfg.bucket}/o?maxResults=1",
|
||||
headers=headers,
|
||||
) as response:
|
||||
if response.status == 403:
|
||||
if response.status == HTTP_FORBIDDEN:
|
||||
msg = f"Access denied to bucket '{cfg.bucket}'. Check permissions."
|
||||
log_structured_entry(
|
||||
"GCS bucket validation failed - access denied - service may not work correctly",
|
||||
(
|
||||
"GCS bucket validation failed - access denied - "
|
||||
"service may not work correctly"
|
||||
),
|
||||
"WARNING",
|
||||
{"bucket": cfg.bucket, "status": response.status}
|
||||
{"bucket": cfg.bucket, "status": response.status},
|
||||
)
|
||||
return msg
|
||||
elif response.status == 404:
|
||||
if response.status == HTTP_NOT_FOUND:
|
||||
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",
|
||||
(
|
||||
"GCS bucket validation failed - not found - "
|
||||
"service may not work correctly"
|
||||
),
|
||||
"WARNING",
|
||||
{"bucket": cfg.bucket, "status": response.status}
|
||||
{"bucket": cfg.bucket, "status": response.status},
|
||||
)
|
||||
return msg
|
||||
elif not response.ok:
|
||||
if 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}
|
||||
{"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(
|
||||
"GCS bucket validation successful", "INFO", {"bucket": cfg.bucket}
|
||||
)
|
||||
return None
|
||||
except Exception as e: # noqa: BLE001
|
||||
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}
|
||||
{"error": str(e), "error_type": type(e).__name__, "bucket": cfg.bucket},
|
||||
)
|
||||
return f"GCS: {str(e)}"
|
||||
return f"GCS: {e!s}"
|
||||
|
||||
|
||||
async def validate_vector_search_access(vs: GoogleCloudVectorSearch, cfg: Settings) -> str | None:
|
||||
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}
|
||||
{"endpoint_name": cfg.endpoint_name},
|
||||
)
|
||||
try:
|
||||
headers = await vs._async_get_auth_headers()
|
||||
session = vs._get_aio_session()
|
||||
headers = await vs._async_get_auth_headers() # noqa: SLF001
|
||||
session = vs._get_aio_session() # noqa: SLF001
|
||||
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."
|
||||
if response.status == HTTP_FORBIDDEN:
|
||||
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",
|
||||
(
|
||||
"Vector search endpoint validation failed - "
|
||||
"access denied - service may not work correctly"
|
||||
),
|
||||
"WARNING",
|
||||
{"endpoint": cfg.endpoint_name, "status": response.status}
|
||||
{"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."
|
||||
if response.status == HTTP_NOT_FOUND:
|
||||
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",
|
||||
(
|
||||
"Vector search endpoint validation failed - "
|
||||
"not found - service may not work correctly"
|
||||
),
|
||||
"WARNING",
|
||||
{"endpoint": cfg.endpoint_name, "status": response.status}
|
||||
{"endpoint": cfg.endpoint_name, "status": response.status},
|
||||
)
|
||||
return msg
|
||||
elif not response.ok:
|
||||
if not response.ok:
|
||||
body = await response.text()
|
||||
msg = f"Failed to access endpoint '{cfg.endpoint_name}': {response.status}"
|
||||
msg = (
|
||||
f"Failed to access endpoint '{cfg.endpoint_name}': "
|
||||
f"{response.status}"
|
||||
)
|
||||
log_structured_entry(
|
||||
"Vector search endpoint validation failed - service may not work correctly",
|
||||
(
|
||||
"Vector search endpoint validation failed - "
|
||||
"service may not work correctly"
|
||||
),
|
||||
"WARNING",
|
||||
{"endpoint": cfg.endpoint_name, "status": response.status, "response": body}
|
||||
{
|
||||
"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(
|
||||
"Vector search endpoint validation successful",
|
||||
"INFO",
|
||||
{"endpoint": cfg.endpoint_name},
|
||||
)
|
||||
return None
|
||||
except Exception as e: # noqa: BLE001
|
||||
log_structured_entry(
|
||||
"Failed to validate vector search endpoint access - service may not work correctly",
|
||||
(
|
||||
"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}
|
||||
{
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"endpoint": cfg.endpoint_name,
|
||||
},
|
||||
)
|
||||
return f"Vector Search: {str(e)}"
|
||||
return f"Vector Search: {e!s}"
|
||||
|
||||
Reference in New Issue
Block a user