feat(semantic_search): make embedding model configurable (#22)

google/gemini-embedding-001 became paid on OpenRouter, leaving
users without credit unable to run the semantic_search tool. The
old code hardcoded the model and 3072 dimensions in three places.

Make OpenRouterEmbedder accept model/dimension via constructor
args or OPENROUTER_EMBEDDING_MODEL / OPENROUTER_EMBEDDING_DIMENSION
env vars, with the previous values as backward-compatible defaults.
Switch the VectorStore and the response payload in mcp_server_main
to read embedder.dimension instead of the hardcoded 3072 so a
configured non-Gemini model does not produce shape mismatches.

Bad dimension input (non-int or non-positive) now raises a clear
ValueError instead of a downstream shape error.

Documented the new env vars in .env.example.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
saidsurucu
2026-05-03 01:46:33 +03:00
co-authored by Claude Opus 4.7
parent ae5d590cca
commit 42731a2c03
3 changed files with 60 additions and 13 deletions
+7
View File
@@ -79,6 +79,13 @@ JWT_SECRET_KEY=your_jwt_secret_key_here
# If not set, semantic search tool will be disabled # If not set, semantic search tool will be disabled
OPENROUTER_API_KEY=sk-or-v1-your_openrouter_api_key_here OPENROUTER_API_KEY=sk-or-v1-your_openrouter_api_key_here
# Optional: override the embedding model and dimension.
# Defaults: google/gemini-embedding-001 at 3072 dims (paid on OpenRouter).
# Pick any embedding model from https://openrouter.ai/models?modality=embedding
# and set the dimension to that model's output size — they must match.
# OPENROUTER_EMBEDDING_MODEL=google/gemini-embedding-001
# OPENROUTER_EMBEDDING_DIMENSION=3072
# ============================================================================= # =============================================================================
# USAGE INSTRUCTIONS # USAGE INSTRUCTIONS
# ============================================================================= # =============================================================================
+3 -2
View File
@@ -1281,7 +1281,7 @@ YANLIŞ KULLANIM:
try: try:
# Initialize components # Initialize components
embedder = OpenRouterEmbedder() embedder = OpenRouterEmbedder()
vector_store = VectorStore(dimension=3072) # Gemini embedding dimension vector_store = VectorStore(dimension=embedder.dimension)
processor = DocumentProcessor(chunk_size=1500, chunk_overlap=300) processor = DocumentProcessor(chunk_size=1500, chunk_overlap=300)
# Step 1: Initial keyword search to get document IDs # Step 1: Initial keyword search to get document IDs
@@ -1437,7 +1437,8 @@ YANLIŞ KULLANIM:
"query": query, "query": query,
"initial_keyword": initial_keyword, "initial_keyword": initial_keyword,
"total_documents_processed": len(documents_data), "total_documents_processed": len(documents_data),
"embedding_dimension": 3072, "embedding_model": embedder.model,
"embedding_dimension": embedder.dimension,
"results": formatted_results, "results": formatted_results,
"stats": { "stats": {
"documents_in_store": stats["num_documents"], "documents_in_store": stats["num_documents"],
+50 -11
View File
@@ -8,6 +8,10 @@ import numpy as np
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
DEFAULT_MODEL = "google/gemini-embedding-001"
DEFAULT_DIMENSION = 3072
def is_openrouter_available() -> bool: def is_openrouter_available() -> bool:
"""Check if OpenRouter API key is available.""" """Check if OpenRouter API key is available."""
return bool(os.getenv("OPENROUTER_API_KEY")) return bool(os.getenv("OPENROUTER_API_KEY"))
@@ -15,16 +19,36 @@ def is_openrouter_available() -> bool:
class OpenRouterEmbedder: class OpenRouterEmbedder:
""" """
Embedder using OpenRouter API with Google's Gemini Embedding model. Embedder using OpenRouter's embedding API.
Requires OPENROUTER_API_KEY environment variable.
The model and dimension are configurable so users can pick any OpenRouter
embedding model (e.g. when one becomes paid or when a different model fits
the budget better). Configuration precedence: explicit constructor args >
environment variables > defaults.
Environment variables:
OPENROUTER_API_KEY (required): OpenRouter credential
OPENROUTER_EMBEDDING_MODEL (optional): override the embedding model id
OPENROUTER_EMBEDDING_DIMENSION (optional): override the vector size
Defaults preserve backward compatibility: ``google/gemini-embedding-001``
at 3072 dimensions.
""" """
def __init__(self): def __init__(self, model: Optional[str] = None, dimension: Optional[int] = None):
""" """
Initialize OpenRouter Embedder. Initialize OpenRouter Embedder.
Args:
model: OpenRouter embedding model id. Falls back to
OPENROUTER_EMBEDDING_MODEL env var, then DEFAULT_MODEL.
dimension: Output vector size. Falls back to
OPENROUTER_EMBEDDING_DIMENSION env var, then DEFAULT_DIMENSION.
Must match the chosen model's actual output size — the vector
store and similarity math rely on it.
Raises: Raises:
ValueError: If OPENROUTER_API_KEY is not set ValueError: If OPENROUTER_API_KEY is not set or dimension is invalid
ImportError: If openai package is not installed ImportError: If openai package is not installed
""" """
api_key = os.getenv("OPENROUTER_API_KEY") api_key = os.getenv("OPENROUTER_API_KEY")
@@ -40,10 +64,25 @@ class OpenRouterEmbedder:
base_url="https://openrouter.ai/api/v1", base_url="https://openrouter.ai/api/v1",
api_key=api_key, api_key=api_key,
) )
self.model = "google/gemini-embedding-001" self.model = model or os.getenv("OPENROUTER_EMBEDDING_MODEL") or DEFAULT_MODEL
self.dimension = 3072
logger.info(f"OpenRouter Embedder initialized with model: {self.model}") dim_value = dimension if dimension is not None else os.getenv("OPENROUTER_EMBEDDING_DIMENSION")
if dim_value is None:
self.dimension = DEFAULT_DIMENSION
else:
try:
self.dimension = int(dim_value)
except (TypeError, ValueError) as e:
raise ValueError(
f"OPENROUTER_EMBEDDING_DIMENSION must be an integer, got {dim_value!r}"
) from e
if self.dimension <= 0:
raise ValueError(f"Embedding dimension must be positive, got {self.dimension}")
logger.info(
f"OpenRouter Embedder initialized with model: {self.model} "
f"(dimension={self.dimension})"
)
def encode_query(self, query: str, task: str = "search result") -> np.ndarray: def encode_query(self, query: str, task: str = "search result") -> np.ndarray:
""" """
@@ -54,7 +93,7 @@ class OpenRouterEmbedder:
task: Task type for prompt template task: Task type for prompt template
Returns: Returns:
Numpy array of embeddings (3072 dimensions) Numpy array of embeddings (``self.dimension`` elements).
""" """
# Apply query prompt template # Apply query prompt template
text = f"task: {task} | query: {query}" text = f"task: {task} | query: {query}"
@@ -93,7 +132,7 @@ class OpenRouterEmbedder:
titles: Optional list of document titles titles: Optional list of document titles
Returns: Returns:
Numpy array of embeddings (N x 3072 dimensions) Numpy array of embeddings (N x ``self.dimension``).
""" """
if not documents: if not documents:
return np.array([]) return np.array([])
@@ -138,8 +177,8 @@ class OpenRouterEmbedder:
Compute cosine similarity between query and documents. Compute cosine similarity between query and documents.
Args: Args:
query_embedding: Query embedding (3072,) query_embedding: Query embedding (``self.dimension``,)
document_embeddings: Document embeddings (N x 3072) document_embeddings: Document embeddings (N x ``self.dimension``)
Returns: Returns:
Similarity scores (N,) Similarity scores (N,)