From 77e2748adedbf731d2456f543fb9e1d9a5ae58d9 Mon Sep 17 00:00:00 2001 From: saidsurucu Date: Sat, 13 Dec 2025 18:06:57 +0300 Subject: [PATCH] feat(semantic-search): Replace local embedding model with OpenRouter API - Replace EmbeddingGemma local model with OpenRouter API integration - Use google/gemini-embedding-001 model via OpenRouter (3072 dimensions) - Add conditional tool registration: auto-disable if OPENROUTER_API_KEY not set - Add openai and numpy dependencies to pyproject.toml - Update .env.example with OPENROUTER_API_KEY configuration - Fix ruff lint issues in semantic_search module --- .env.example | 9 + mcp_server_main.py | 508 +++++++++++++++++------------------- pyproject.toml | 4 +- semantic_search/__init__.py | 4 +- semantic_search/embedder.py | 219 +++++++--------- 5 files changed, 356 insertions(+), 388 deletions(-) diff --git a/.env.example b/.env.example index 270a0fb..90fbd3e 100644 --- a/.env.example +++ b/.env.example @@ -70,6 +70,15 @@ JWT_SECRET_KEY=your_jwt_secret_key_here # MAX_REQUESTS_PER_MINUTE=60 # BURST_CAPACITY=20 +# ============================================================================= +# SEMANTIC SEARCH SETTINGS (Optional) +# ============================================================================= + +# OpenRouter API Key for semantic search functionality +# Get your API key from: https://openrouter.ai/keys +# If not set, semantic search tool will be disabled +OPENROUTER_API_KEY=sk-or-v1-your_openrouter_api_key_here + # ============================================================================= # USAGE INSTRUCTIONS # ============================================================================= diff --git a/mcp_server_main.py b/mcp_server_main.py index d78dcfa..3e544cb 100644 --- a/mcp_server_main.py +++ b/mcp_server_main.py @@ -2,14 +2,12 @@ import asyncio import atexit import logging -import os import httpx import json import time from collections import defaultdict -from pydantic import BaseModel, HttpUrl, Field -from typing import Optional, Dict, List, Literal, Any, Union -import urllib.parse +from pydantic import HttpUrl, Field +from typing import Optional, Dict, List, Literal, Any from fastmcp.server.middleware import Middleware, MiddlewareContext # Optional tiktoken import for token counting @@ -150,7 +148,7 @@ class TokenCountingMiddleware(Middleware): return result - except Exception as e: + except Exception: duration_ms = (time.perf_counter() - start_time) * 1000 self.log_token_usage("tool_call_error", input_tokens, 0, tool_name, duration_ms) @@ -180,7 +178,7 @@ class TokenCountingMiddleware(Middleware): return result - except Exception as e: + except Exception: duration_ms = (time.perf_counter() - start_time) * 1000 self.log_token_usage("resource_read_error", 0, 0, resource_uri, duration_ms) @@ -210,7 +208,7 @@ class TokenCountingMiddleware(Middleware): return result - except Exception as e: + except Exception: duration_ms = (time.perf_counter() - start_time) * 1000 self.log_token_usage("prompt_get_error", 0, 0, prompt_name, duration_ms) @@ -251,10 +249,6 @@ def create_app(auth=None): # --- Module Imports --- from yargitay_mcp_module.client import YargitayOfficialApiClient -from yargitay_mcp_module.models import ( - YargitayDetailedSearchRequest, YargitayDocumentMarkdown, CompactYargitaySearchResult, - YargitayBirimEnum, CleanYargitayDecisionEntry -) from bedesten_mcp_module.client import BedestenApiClient from bedesten_mcp_module.models import ( BedestenSearchRequest, BedestenSearchData, @@ -262,66 +256,49 @@ from bedesten_mcp_module.models import ( ) from bedesten_mcp_module.enums import BirimAdiEnum -# Semantic Search Module Imports -from semantic_search.embedder import EmbeddingGemma -from semantic_search.vector_store import VectorStore -from semantic_search.processor import DocumentProcessor +# Semantic Search Module Imports (conditional based on OPENROUTER_API_KEY) +from semantic_search.embedder import is_openrouter_available +SEMANTIC_SEARCH_AVAILABLE = is_openrouter_available() + +if SEMANTIC_SEARCH_AVAILABLE: + from semantic_search.embedder import OpenRouterEmbedder + from semantic_search.vector_store import VectorStore + from semantic_search.processor import DocumentProcessor + logger.info("Semantic search enabled (OPENROUTER_API_KEY found)") +else: + logger.info("Semantic search disabled (OPENROUTER_API_KEY not set)") from danistay_mcp_module.client import DanistayApiClient -from danistay_mcp_module.models import ( - DanistayKeywordSearchRequest, DanistayDetailedSearchRequest, - DanistayDocumentMarkdown, CompactDanistaySearchResult -) from emsal_mcp_module.client import EmsalApiClient from emsal_mcp_module.models import ( - EmsalSearchRequest, EmsalDocumentMarkdown, CompactEmsalSearchResult + EmsalSearchRequest, CompactEmsalSearchResult ) from uyusmazlik_mcp_module.client import UyusmazlikApiClient from uyusmazlik_mcp_module.models import ( - UyusmazlikSearchRequest, UyusmazlikSearchResponse, UyusmazlikDocumentMarkdown, - UyusmazlikBolumEnum, UyusmazlikTuruEnum, UyusmazlikKararSonucuEnum + UyusmazlikSearchRequest, UyusmazlikBolumEnum, UyusmazlikTuruEnum, UyusmazlikKararSonucuEnum ) from anayasa_mcp_module.client import AnayasaMahkemesiApiClient from anayasa_mcp_module.bireysel_client import AnayasaBireyselBasvuruApiClient from anayasa_mcp_module.unified_client import AnayasaUnifiedClient from anayasa_mcp_module.models import ( - AnayasaNormDenetimiSearchRequest, - AnayasaSearchResult, - AnayasaDocumentMarkdown, - AnayasaBireyselReportSearchRequest, - AnayasaBireyselReportSearchResult, - AnayasaBireyselBasvuruDocumentMarkdown, AnayasaUnifiedSearchRequest, - AnayasaUnifiedSearchResult, - AnayasaUnifiedDocumentMarkdown, # Removed enum imports - now using Literal strings in models ) # KIK v2 Module Imports (New API) from kik_mcp_module.client_v2 import KikV2ApiClient from kik_mcp_module.models_v2 import KikV2DecisionType -from kik_mcp_module.models_v2 import ( - KikV2SearchResult, - KikV2DocumentMarkdown -) from rekabet_mcp_module.client import RekabetKurumuApiClient from rekabet_mcp_module.models import ( RekabetKurumuSearchRequest, RekabetSearchResult, - RekabetDocument, RekabetKararTuruGuidEnum ) from sayistay_mcp_module.client import SayistayApiClient from sayistay_mcp_module.models import ( - GenelKurulSearchRequest, GenelKurulSearchResponse, - TemyizKuruluSearchRequest, TemyizKuruluSearchResponse, - DaireSearchRequest, DaireSearchResponse, - SayistayDocumentMarkdown, - SayistayUnifiedSearchRequest, SayistayUnifiedSearchResult, - SayistayUnifiedDocumentMarkdown + SayistayUnifiedSearchRequest ) -from sayistay_mcp_module.enums import DaireEnum, KamuIdaresiTuruEnum, WebKararKonusuEnum from sayistay_mcp_module.unified_client import SayistayUnifiedClient # KVKK Module Imports @@ -335,14 +312,11 @@ from kvkk_mcp_module.models import ( # BDDK Module Imports from bddk_mcp_module.client import BddkApiClient from bddk_mcp_module.models import ( - BddkSearchRequest, - BddkSearchResult, - BddkDocumentMarkdown + BddkSearchRequest ) # Create a placeholder app that will be properly initialized after tools are defined -from fastmcp import FastMCP # MCP app for Turkish legal databases with explicit capabilities app = FastMCP( @@ -652,7 +626,7 @@ async def search_emsal_detailed_decisions( page_size=page_size ) - logger.info(f"Tool 'search_emsal_detailed_decisions' called.") + logger.info("Tool 'search_emsal_detailed_decisions' called.") try: api_response = await emsal_client_instance.search_detailed_decisions(search_query) if api_response.data: @@ -664,8 +638,8 @@ async def search_emsal_detailed_decisions( ).model_dump() logger.warning("API response for Emsal search did not contain expected data structure.") return CompactEmsalSearchResult(decisions=[], total_records=0, requested_page=search_query.page_number, page_size=search_query.page_size).model_dump() - except Exception as e: - logger.exception(f"Error in tool 'search_emsal_detailed_decisions'.") + except Exception: + logger.exception("Error in tool 'search_emsal_detailed_decisions'.") raise @app.tool( @@ -682,8 +656,8 @@ async def get_emsal_document_markdown(id: str) -> Dict[str, Any]: try: result = await emsal_client_instance.get_decision_document_as_markdown(id) return result.model_dump() - except Exception as e: - logger.exception(f"Error in tool 'get_emsal_document_markdown'.") + except Exception: + logger.exception("Error in tool 'get_emsal_document_markdown'.") raise # --- MCP Tools for Uyusmazlik --- @@ -751,12 +725,12 @@ async def search_uyusmazlik_decisions( not_hepsi=not_hepsi ) - logger.info(f"Tool 'search_uyusmazlik_decisions' called.") + logger.info("Tool 'search_uyusmazlik_decisions' called.") try: result = await uyusmazlik_client_instance.search_decisions(search_params) return result.model_dump() - except Exception as e: - logger.exception(f"Error in tool 'search_uyusmazlik_decisions'.") + except Exception: + logger.exception("Error in tool 'search_uyusmazlik_decisions'.") raise @app.tool( @@ -776,8 +750,8 @@ async def get_uyusmazlik_document_markdown_from_url( try: result = await uyusmazlik_client_instance.get_decision_document_as_markdown(str(document_url)) return result.model_dump() - except Exception as e: - logger.exception(f"Error in tool 'get_uyusmazlik_document_markdown_from_url'.") + except Exception: + logger.exception("Error in tool 'get_uyusmazlik_document_markdown_from_url'.") raise # --- DEACTIVATED: MCP Tools for Anayasa Mahkemesi (Individual Tools) --- @@ -868,8 +842,8 @@ async def search_anayasa_unified( result = await anayasa_unified_client_instance.search_unified(request) return json.dumps(result.model_dump(), ensure_ascii=False, indent=2) - except Exception as e: - logger.exception(f"Error in tool 'search_anayasa_unified'.") + except Exception: + logger.exception("Error in tool 'search_anayasa_unified'.") raise @app.tool( @@ -890,8 +864,8 @@ async def get_anayasa_document_unified( result = await anayasa_unified_client_instance.get_document_unified(document_url, page_number) return json.dumps(result.model_dump(mode='json'), ensure_ascii=False, indent=2) - except Exception as e: - logger.exception(f"Error in tool 'get_anayasa_document_unified'.") + except Exception: + logger.exception("Error in tool 'get_anayasa_document_unified'.") raise # --- MCP Tools for KIK v2 (Kamu İhale Kurulu - New API) --- @@ -1066,7 +1040,7 @@ async def search_rekabet_kurumu_decisions( result = await rekabet_client_instance.search_decisions(search_query) return result.model_dump() - except Exception as e: + except Exception: logger.exception("Error in tool 'search_rekabet_kurumu_decisions'.") return RekabetSearchResult(decisions=[], retrieved_page_number=page, total_records_found=0, total_pages=0).model_dump() @@ -1089,7 +1063,7 @@ async def get_rekabet_kurumu_document( try: result = await rekabet_client_instance.get_decision_document(karar_id, page_number=current_page_to_fetch) return result.model_dump() - except Exception as e: + except Exception: logger.exception(f"Error in tool 'get_rekabet_kurumu_document'. Karar ID: {karar_id}") raise @@ -1200,7 +1174,7 @@ For best results, use exact phrases with quotes for legal terms."""), "page_size": pageSize, "searched_courts": court_types } - except Exception as e: + except Exception: logger.exception("Error in tool 'search_bedesten_unified'") raise @@ -1222,224 +1196,226 @@ async def get_bedesten_document_markdown( try: return await bedesten_client_instance.get_document_as_markdown(documentId) - except Exception as e: + except Exception: logger.exception("Error in tool 'get_kyb_bedesten_document_markdown'") raise -# --- Semantic Search Tool --- -@app.tool( - description="Perform semantic search on Turkish legal decisions using EmbeddingGemma for intelligent re-ranking", - annotations={ - "readOnlyHint": True, - "openWorldHint": True, - "idempotentHint": True - } -) -async def search_bedesten_semantic( - query: str = Field(..., description="Search query in Turkish for semantic matching"), - initial_keyword: str = Field(..., description="Initial keyword for Bedesten API search (broad term)"), - court_types: List[BedestenCourtTypeEnum] = Field( - default=["YARGITAYKARARI", "DANISTAYKARAR", "YERELHUKUK", "ISTINAFHUKUK", "KYB"], - description="Court types to search: YARGITAYKARARI, DANISTAYKARAR, YERELHUKUK, ISTINAFHUKUK, KYB (default: all)" - ), - top_k: int = Field(10, ge=1, le=50, description="Number of top results to return (1-50)") -) -> Dict[str, Any]: - """ - Perform semantic search on Turkish legal decisions using EmbeddingGemma. +# --- Semantic Search Tool (Conditional - requires OPENROUTER_API_KEY) --- +if SEMANTIC_SEARCH_AVAILABLE: + @app.tool( + description="Perform semantic search on Turkish legal decisions using OpenRouter Gemini embeddings for intelligent re-ranking", + annotations={ + "readOnlyHint": True, + "openWorldHint": True, + "idempotentHint": True + } + ) + async def search_bedesten_semantic( + query: str = Field(..., description="Search query in Turkish for semantic matching"), + initial_keyword: str = Field(..., description="Initial keyword for Bedesten API search (broad term)"), + court_types: List[BedestenCourtTypeEnum] = Field( + default=["YARGITAYKARARI", "DANISTAYKARAR", "YERELHUKUK", "ISTINAFHUKUK", "KYB"], + description="Court types to search: YARGITAYKARARI, DANISTAYKARAR, YERELHUKUK, ISTINAFHUKUK, KYB (default: all)" + ), + top_k: int = Field(10, ge=1, le=50, description="Number of top results to return (1-50)") + ) -> Dict[str, Any]: + """ + Perform semantic search on Turkish legal decisions using OpenRouter API. - This tool: - 1. Searches Bedesten API with initial keyword (retrieves 100 results) - 2. Fetches full document content for each result - 3. Generates embeddings using Google's EmbeddingGemma model - 4. Performs semantic similarity search with the query - 5. Returns re-ranked results based on semantic relevance + This tool: + 1. Searches Bedesten API with initial keyword (retrieves 100 results) + 2. Fetches full document content for each result + 3. Generates embeddings using Google's Gemini Embedding model via OpenRouter + 4. Performs semantic similarity search with the query + 5. Returns re-ranked results based on semantic relevance - Benefits over keyword search: - - Better understanding of context and meaning - - Finds semantically similar documents even with different wording - - More accurate ranking based on relevance - - Supports multilingual queries (100+ languages) - """ - logger.info(f"Semantic search tool called with query: {query}, keyword: {initial_keyword}") + Benefits over keyword search: + - Better understanding of context and meaning + - Finds semantically similar documents even with different wording + - More accurate ranking based on relevance + - Supports multilingual queries (100+ languages) - try: - # Initialize components - embedder = EmbeddingGemma() - vector_store = VectorStore(dimension=256) - processor = DocumentProcessor(chunk_size=1500, chunk_overlap=300) + Note: Requires OPENROUTER_API_KEY environment variable to be set. + """ + logger.info(f"Semantic search tool called with query: {query}, keyword: {initial_keyword}") - # Step 1: Initial keyword search to get document IDs - logger.info(f"Step 1: Searching Bedesten API with keyword: {initial_keyword}") + try: + # Initialize components + embedder = OpenRouterEmbedder() + vector_store = VectorStore(dimension=3072) # Gemini embedding dimension + processor = DocumentProcessor(chunk_size=1500, chunk_overlap=300) - all_decisions = [] + # Step 1: Initial keyword search to get document IDs + logger.info(f"Step 1: Searching Bedesten API with keyword: {initial_keyword}") - # Search each court type - for court_type in court_types: - try: - per_court_limit = max(20, 100 // len(court_types)) + all_decisions = [] - search_results = await bedesten_client_instance.search_documents( - BedestenSearchRequest( - data=BedestenSearchData( - phrase=initial_keyword, - itemTypeList=[court_type], - pageSize=per_court_limit, - pageNumber=1 + # Search each court type + for court_type in court_types: + try: + per_court_limit = max(20, 100 // len(court_types)) + + search_results = await bedesten_client_instance.search_documents( + BedestenSearchRequest( + data=BedestenSearchData( + phrase=initial_keyword, + itemTypeList=[court_type], + pageSize=per_court_limit, + pageNumber=1 + ) ) ) - ) - if search_results.data and search_results.data.emsalKararList: - all_decisions.extend(search_results.data.emsalKararList) - logger.info(f"Found {len(search_results.data.emsalKararList)} results from {court_type}") + if search_results.data and search_results.data.emsalKararList: + all_decisions.extend(search_results.data.emsalKararList) + logger.info(f"Found {len(search_results.data.emsalKararList)} results from {court_type}") - except Exception as e: - logger.warning(f"Error searching {court_type}: {e}") + except Exception as e: + logger.warning(f"Error searching {court_type}: {e}") + + if not all_decisions: + logger.warning("No documents found from initial search") + return { + "status": "no_results", + "message": "No documents found matching the initial keyword", + "results": [] + } + + logger.info(f"Total documents found: {len(all_decisions)}") + + # Step 2: Fetch document content and process + logger.info("Step 2: Fetching and processing document content...") + + documents_data = [] + failed_fetches = 0 + decisions_to_process = all_decisions[:100] + + for i, decision in enumerate(decisions_to_process): + try: + doc = await bedesten_client_instance.get_document_as_markdown(decision.documentId) + + if doc.markdown_content: + metadata = { + "document_id": decision.documentId, + "birim_adi": decision.birimAdi, + "esas_no": decision.esasNo, + "karar_no": decision.kararNo, + "karar_tarihi": decision.kararTarihiStr, + "court_type": decision.itemType.name if decision.itemType else None + } + + chunks = processor.process_document( + document_id=decision.documentId, + text=doc.markdown_content, + metadata=metadata + ) + + if chunks: + full_text = " ".join([chunk.text for chunk in chunks]) + documents_data.append({ + "id": decision.documentId, + "text": full_text[:3000], + "metadata": metadata + }) + + if (i + 1) % 10 == 0: + logger.info(f"Processed {i + 1}/{len(decisions_to_process)} documents") + + except Exception as e: + logger.warning(f"Failed to fetch document {decision.documentId}: {e}") + failed_fetches += 1 + + if not documents_data: + logger.warning("No documents could be processed") + return { + "status": "processing_error", + "message": "Could not process any documents", + "results": [] + } + + logger.info(f"Successfully processed {len(documents_data)} documents, {failed_fetches} failed") + + # Step 3: Generate embeddings + logger.info("Step 3: Generating embeddings...") + + query_embedding = embedder.encode_query(query, task="search result") + + doc_texts = [doc["text"] for doc in documents_data] + doc_titles = [doc["metadata"].get("birim_adi", "none") for doc in documents_data] + doc_embeddings = embedder.encode_documents(doc_texts, titles=doc_titles) + + # No dimension reduction - using full 3072 dimensions + + # Step 4: Add to vector store and search + logger.info("Step 4: Performing semantic search...") + + doc_ids = [doc["id"] for doc in documents_data] + doc_metadatas = [doc["metadata"] for doc in documents_data] + + vector_store.add_documents( + ids=doc_ids, + texts=doc_texts, + embeddings=doc_embeddings, + metadata=doc_metadatas + ) + + search_results = vector_store.search( + query_embedding=query_embedding, + top_k=top_k, + threshold=0.3 + ) + + # Step 5: Format results + logger.info(f"Step 5: Formatting {len(search_results)} results") + + formatted_results = [] + for doc, score in search_results: + title_parts = [] + if doc.metadata.get("birim_adi"): + title_parts.append(doc.metadata["birim_adi"]) + if doc.metadata.get("esas_no"): + title_parts.append(f"Esas: {doc.metadata['esas_no']}") + if doc.metadata.get("karar_no"): + title_parts.append(f"Karar: {doc.metadata['karar_no']}") + if doc.metadata.get("karar_tarihi"): + title_parts.append(f"Tarih: {doc.metadata['karar_tarihi']}") + + title = " - ".join(title_parts) if title_parts else f"Document {doc.id}" + + formatted_results.append({ + "document_id": doc.id, + "title": title, + "similarity_score": float(score), + "preview": doc.text[:500] + "..." if len(doc.text) > 500 else doc.text, + "metadata": doc.metadata, + "source_url": f"https://mevzuat.adalet.gov.tr/ictihat/{doc.id}" + }) + + stats = vector_store.get_stats() - if not all_decisions: - logger.warning("No documents found from initial search") return { - "status": "no_results", - "message": "No documents found matching the initial keyword", + "status": "success", + "query": query, + "initial_keyword": initial_keyword, + "total_documents_processed": len(documents_data), + "embedding_dimension": 3072, + "results": formatted_results, + "stats": { + "documents_in_store": stats["num_documents"], + "memory_usage_mb": round(stats["memory_usage_mb"], 2), + "failed_fetches": failed_fetches + } + } + + except Exception as e: + logger.exception(f"Error in semantic search: {e}") + return { + "status": "error", + "message": str(e), "results": [] } - logger.info(f"Total documents found: {len(all_decisions)}") - - # Step 2: Fetch document content and process - logger.info("Step 2: Fetching and processing document content...") - - documents_data = [] - failed_fetches = 0 - decisions_to_process = all_decisions[:100] - - for i, decision in enumerate(decisions_to_process): - try: - doc = await bedesten_client_instance.get_document_as_markdown(decision.documentId) - - if doc.markdown_content: - metadata = { - "document_id": decision.documentId, - "birim_adi": decision.birimAdi, - "esas_no": decision.esasNo, - "karar_no": decision.kararNo, - "karar_tarihi": decision.kararTarihiStr, - "court_type": decision.itemType.name if decision.itemType else None - } - - chunks = processor.process_document( - document_id=decision.documentId, - text=doc.markdown_content, - metadata=metadata - ) - - if chunks: - full_text = " ".join([chunk.text for chunk in chunks]) - documents_data.append({ - "id": decision.documentId, - "text": full_text[:3000], - "metadata": metadata - }) - - if (i + 1) % 10 == 0: - logger.info(f"Processed {i + 1}/{len(decisions_to_process)} documents") - - except Exception as e: - logger.warning(f"Failed to fetch document {decision.documentId}: {e}") - failed_fetches += 1 - - if not documents_data: - logger.warning("No documents could be processed") - return { - "status": "processing_error", - "message": "Could not process any documents", - "results": [] - } - - logger.info(f"Successfully processed {len(documents_data)} documents, {failed_fetches} failed") - - # Step 3: Generate embeddings - logger.info("Step 3: Generating embeddings...") - - query_embedding = embedder.encode_query(query, task="search result") - - doc_texts = [doc["text"] for doc in documents_data] - doc_titles = [doc["metadata"].get("birim_adi", "none") for doc in documents_data] - doc_embeddings = embedder.encode_documents(doc_texts, titles=doc_titles) - - query_embedding = embedder.reduce_dimensions(query_embedding, 256) - doc_embeddings = embedder.reduce_dimensions(doc_embeddings, 256) - - # Step 4: Add to vector store and search - logger.info("Step 4: Performing semantic search...") - - doc_ids = [doc["id"] for doc in documents_data] - doc_metadatas = [doc["metadata"] for doc in documents_data] - - vector_store.add_documents( - ids=doc_ids, - texts=doc_texts, - embeddings=doc_embeddings, - metadata=doc_metadatas - ) - - search_results = vector_store.search( - query_embedding=query_embedding, - top_k=top_k, - threshold=0.3 - ) - - # Step 5: Format results - logger.info(f"Step 5: Formatting {len(search_results)} results") - - formatted_results = [] - for doc, score in search_results: - title_parts = [] - if doc.metadata.get("birim_adi"): - title_parts.append(doc.metadata["birim_adi"]) - if doc.metadata.get("esas_no"): - title_parts.append(f"Esas: {doc.metadata['esas_no']}") - if doc.metadata.get("karar_no"): - title_parts.append(f"Karar: {doc.metadata['karar_no']}") - if doc.metadata.get("karar_tarihi"): - title_parts.append(f"Tarih: {doc.metadata['karar_tarihi']}") - - title = " - ".join(title_parts) if title_parts else f"Document {doc.id}" - - formatted_results.append({ - "document_id": doc.id, - "title": title, - "similarity_score": float(score), - "preview": doc.text[:500] + "..." if len(doc.text) > 500 else doc.text, - "metadata": doc.metadata, - "source_url": f"https://mevzuat.adalet.gov.tr/ictihat/{doc.id}" - }) - - stats = vector_store.get_stats() - - return { - "status": "success", - "query": query, - "initial_keyword": initial_keyword, - "total_documents_processed": len(documents_data), - "embedding_dimension": 256, - "results": formatted_results, - "stats": { - "documents_in_store": stats["num_documents"], - "memory_usage_mb": round(stats["memory_usage_mb"], 2), - "failed_fetches": failed_fetches - } - } - - except Exception as e: - logger.exception(f"Error in semantic search: {e}") - return { - "status": "error", - "message": str(e), - "results": [] - } - # --- MCP Tools for Sayıştay (Turkish Court of Accounts) --- @@ -1627,7 +1603,7 @@ async def search_sayistay_unified( ) result = await sayistay_unified_client_instance.search_unified(search_request) return result.model_dump() - except Exception as e: + except Exception: logger.exception("Error in tool 'search_sayistay_unified'") raise @@ -1652,7 +1628,7 @@ async def get_sayistay_document_unified( try: result = await sayistay_unified_client_instance.get_document_unified(decision_id, decision_type) return result.model_dump() - except Exception as e: + except Exception: logger.exception("Error in tool 'get_sayistay_document_unified'") raise @@ -2267,7 +2243,7 @@ async def search( ] } - except Exception as e: + except Exception: logger.exception("Error in ChatGPT Deep Research search tool") # Return partial results if any were found if results: @@ -2406,7 +2382,7 @@ async def fetch( doc = await bedesten_client_instance.get_document_as_markdown(doc_id) """ - except Exception as e: + except Exception: logger.exception(f"Error fetching ChatGPT Deep Research document {id}") raise @@ -2424,7 +2400,7 @@ def main(): app.run() except KeyboardInterrupt: logger.info("Server shut down by user (KeyboardInterrupt).") - except Exception as e: + except Exception: logger.exception("Server failed to start or crashed.") finally: logger.info(f"{app.name} server has shut down.") diff --git a/pyproject.toml b/pyproject.toml index 7596c13..1fd28e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,8 @@ dependencies = [ "pypdf>=5.5.0", "fastapi>=0.115.14", "cryptography>=44.0.0", + "openai>=1.0.0", + "numpy>=1.24.0", ] [project.optional-dependencies] @@ -59,7 +61,7 @@ yargi-mcp = "mcp_server_main:main" py-modules = ["mcp_server_main", "mcp_auth_factory", "mcp_auth_http_adapter", "asgi_app", "fastapi_app", "starlette_app", "run_asgi", "stripe_webhook"] [tool.setuptools.packages.find] -include = ["*_mcp_module", "mcp_auth"] +include = ["*_mcp_module", "mcp_auth", "semantic_search"] [build-system] requires = ["setuptools>=65.0", "wheel"] diff --git a/semantic_search/__init__.py b/semantic_search/__init__.py index c1e9e74..e773900 100644 --- a/semantic_search/__init__.py +++ b/semantic_search/__init__.py @@ -1,7 +1,7 @@ # semantic_search/__init__.py -from .embedder import EmbeddingGemma +from .embedder import OpenRouterEmbedder, is_openrouter_available from .vector_store import VectorStore from .processor import DocumentProcessor -__all__ = ['EmbeddingGemma', 'VectorStore', 'DocumentProcessor'] \ No newline at end of file +__all__ = ['OpenRouterEmbedder', 'is_openrouter_available', 'VectorStore', 'DocumentProcessor'] diff --git a/semantic_search/embedder.py b/semantic_search/embedder.py index 29f9452..6c9c028 100644 --- a/semantic_search/embedder.py +++ b/semantic_search/embedder.py @@ -1,173 +1,154 @@ # semantic_search/embedder.py import logging +import os from typing import List, Optional import numpy as np -from sentence_transformers import SentenceTransformer -import torch logger = logging.getLogger(__name__) -class EmbeddingGemma: + +def is_openrouter_available() -> bool: + """Check if OpenRouter API key is available.""" + return bool(os.getenv("OPENROUTER_API_KEY")) + + +class OpenRouterEmbedder: """ - Wrapper for Google's EmbeddingGemma model. - Handles query and document encoding with proper prompt templates. + Embedder using OpenRouter API with Google's Gemini Embedding model. + Requires OPENROUTER_API_KEY environment variable. """ - - def __init__(self, model_name: str = "google/embeddinggemma-300m", device: Optional[str] = None): + + def __init__(self): """ - Initialize EmbeddingGemma model. - - Args: - model_name: HuggingFace model name - device: Device to run model on ('cuda', 'cpu', or None for auto) + Initialize OpenRouter Embedder. + + Raises: + ValueError: If OPENROUTER_API_KEY is not set + ImportError: If openai package is not installed """ - self.model_name = model_name - - # Auto-detect device if not specified - if device is None: - self.device = 'cuda' if torch.cuda.is_available() else 'cpu' - else: - self.device = device - - logger.info(f"Initializing EmbeddingGemma on device: {self.device}") - + api_key = os.getenv("OPENROUTER_API_KEY") + if not api_key: + raise ValueError("OPENROUTER_API_KEY environment variable is not set") + try: - # Load model with float32 precision (EmbeddingGemma doesn't support float16) - self.model = SentenceTransformer(model_name, device=self.device) - self.model.eval() # Set to evaluation mode - - # Set precision to float32 or bfloat16 - if self.device == 'cuda' and torch.cuda.is_bf16_supported(): - logger.info("Using bfloat16 precision for CUDA") - self.dtype = torch.bfloat16 - else: - logger.info("Using float32 precision") - self.dtype = torch.float32 - - logger.info(f"Successfully loaded model: {model_name}") - - except Exception as e: - logger.error(f"Failed to load EmbeddingGemma model: {e}") - raise - + from openai import OpenAI + except ImportError: + raise ImportError("openai package is required. Install with: pip install openai") + + self.client = OpenAI( + base_url="https://openrouter.ai/api/v1", + api_key=api_key, + ) + self.model = "google/gemini-embedding-001" + self.dimension = 3072 + + logger.info(f"OpenRouter Embedder initialized with model: {self.model}") + def encode_query(self, query: str, task: str = "search result") -> np.ndarray: """ - Encode a search query with appropriate prompt template. - + Encode a search query. + Args: query: The search query text - task: Task type for prompt template (search result, question answering, etc.) - + task: Task type for prompt template + Returns: - Numpy array of embeddings (768 dimensions) + Numpy array of embeddings (3072 dimensions) """ # Apply query prompt template - prompted_query = f"task: {task} | query: {query}" - + text = f"task: {task} | query: {query}" + try: - with torch.no_grad(): - # Encode with model - embeddings = self.model.encode( - prompted_query, - convert_to_numpy=True, - normalize_embeddings=True, # L2 normalization for cosine similarity - show_progress_bar=False - ) - - logger.debug(f"Encoded query: {query[:50]}... -> shape: {embeddings.shape}") - return embeddings - + response = self.client.embeddings.create( + model=self.model, + input=text, + encoding_format="float", + extra_headers={ + "HTTP-Referer": "https://yargimcp.com", + "X-Title": "Yargi MCP Server", + } + ) + + embedding = np.array(response.data[0].embedding, dtype=np.float32) + + # L2 normalize for cosine similarity + norm = np.linalg.norm(embedding) + if norm > 0: + embedding = embedding / norm + + logger.debug(f"Encoded query: {query[:50]}... -> shape: {embedding.shape}") + return embedding + except Exception as e: logger.error(f"Failed to encode query: {e}") raise - + def encode_documents(self, documents: List[str], titles: Optional[List[str]] = None) -> np.ndarray: """ - Encode multiple documents with appropriate prompt template. - + Encode multiple documents with batch API call. + Args: documents: List of document texts titles: Optional list of document titles - + Returns: - Numpy array of embeddings (N x 768 dimensions) + Numpy array of embeddings (N x 3072 dimensions) """ if not documents: return np.array([]) - + # Apply document prompt template - prompted_docs = [] + texts = [] for i, doc in enumerate(documents): title = titles[i] if titles and i < len(titles) else "none" - prompted_doc = f"title: {title} | text: {doc}" - prompted_docs.append(prompted_doc) - + text = f"title: {title} | text: {doc}" + texts.append(text) + try: - with torch.no_grad(): - # Batch encode documents - embeddings = self.model.encode( - prompted_docs, - convert_to_numpy=True, - normalize_embeddings=True, - show_progress_bar=len(documents) > 10, - batch_size=8 # Adjust based on memory - ) - + response = self.client.embeddings.create( + model=self.model, + input=texts, + encoding_format="float", + extra_headers={ + "HTTP-Referer": "https://yargimcp.com", + "X-Title": "Yargi MCP Server", + } + ) + + # Extract embeddings in order + embeddings = np.array( + [d.embedding for d in sorted(response.data, key=lambda x: x.index)], + dtype=np.float32 + ) + + # L2 normalize each embedding for cosine similarity + norms = np.linalg.norm(embeddings, axis=1, keepdims=True) + embeddings = embeddings / (norms + 1e-8) + logger.info(f"Encoded {len(documents)} documents -> shape: {embeddings.shape}") return embeddings - + except Exception as e: logger.error(f"Failed to encode documents: {e}") raise - - def reduce_dimensions(self, embeddings: np.ndarray, target_dim: int = 512) -> np.ndarray: - """ - Reduce embedding dimensions using Matryoshka Representation Learning. - - Args: - embeddings: Original embeddings (N x 768) - target_dim: Target dimension (512, 256, or 128) - - Returns: - Reduced embeddings (N x target_dim) - """ - if target_dim not in [512, 256, 128]: - raise ValueError(f"Target dimension must be 512, 256, or 128, got {target_dim}") - - if len(embeddings.shape) == 1: - # Single embedding - reduced = embeddings[:target_dim] - # Re-normalize after truncation - norm = np.linalg.norm(reduced) - if norm > 0: - reduced = reduced / norm - else: - # Multiple embeddings - reduced = embeddings[:, :target_dim] - # Re-normalize each embedding - norms = np.linalg.norm(reduced, axis=1, keepdims=True) - reduced = reduced / (norms + 1e-8) # Avoid division by zero - - logger.debug(f"Reduced dimensions: {embeddings.shape} -> {reduced.shape}") - return reduced - + def compute_similarity(self, query_embedding: np.ndarray, document_embeddings: np.ndarray) -> np.ndarray: """ Compute cosine similarity between query and documents. - + Args: - query_embedding: Query embedding (768,) - document_embeddings: Document embeddings (N x 768) - + query_embedding: Query embedding (3072,) + document_embeddings: Document embeddings (N x 3072) + Returns: Similarity scores (N,) """ # Ensure query is 2D for matrix multiplication if len(query_embedding.shape) == 1: query_embedding = query_embedding.reshape(1, -1) - + # Compute cosine similarity (embeddings are already normalized) similarities = np.dot(document_embeddings, query_embedding.T).squeeze() - - return similarities \ No newline at end of file + + return similarities