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:
co-authored by
Claude Opus 4.7
parent
ae5d590cca
commit
42731a2c03
@@ -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
@@ -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
@@ -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,)
|
||||||
|
|||||||
Reference in New Issue
Block a user