@iflow-mcp/apple-rag-mcp 4.6.2

package/src/services/rate-limit.ts

@@ -0,0 +1,141 @@
+/**
+ * Rate Limiting Service with D1 timeout protection
+ */
+
+import type { AuthContext } from "../types/index.js";
+import { withD1Timeout } from "../utils/d1-utils.js";
+import { logger } from "../utils/logger.js";
+
+interface RateLimitResult {
+  allowed: boolean;
+  limit: number;
+  remaining: number;
+  resetAt: string;
+  planType: string;
+  limitType: "weekly" | "minute";
+  minuteLimit?: number;
+  minuteRemaining?: number;
+  minuteResetAt?: string;
+}
+
+interface PlanLimits {
+  weeklyQueries: number;
+  requestsPerMinute: number;
+}
+
+const PLAN_LIMITS: Record<string, PlanLimits> = {
+  anonymous: { weeklyQueries: 30, requestsPerMinute: 3 },
+  hobby: { weeklyQueries: 50, requestsPerMinute: 5 },
+  pro: { weeklyQueries: 50000, requestsPerMinute: 50 },
+  enterprise: { weeklyQueries: -1, requestsPerMinute: -1 },
+};
+
+export class RateLimitService {
+  constructor(private d1: D1Database) {}
+
+  async checkLimits(
+    clientIP: string,
+    authContext: AuthContext
+  ): Promise<RateLimitResult> {
+    const identifier = authContext.userId || `anon_${clientIP}`;
+    const planType = authContext.isAuthenticated && authContext.userId
+      ? await this.getPlanType(authContext.userId)
+      : "anonymous";
+
+    const limits = PLAN_LIMITS[planType] || PLAN_LIMITS.hobby;
+
+    const [weeklyUsage, minuteUsage] = await Promise.all([
+      this.getUsageCount(identifier, "weekly"),
+      this.getUsageCount(identifier, "minute"),
+    ]);
+
+    const weeklyAllowed = limits.weeklyQueries === -1 || weeklyUsage < limits.weeklyQueries;
+    const minuteAllowed = limits.requestsPerMinute === -1 || minuteUsage < limits.requestsPerMinute;
+    const allowed = weeklyAllowed && minuteAllowed;
+
+    if (!allowed) {
+      logger.info(
+        `Rate limit: ${identifier} (${planType}) weekly=${weeklyUsage}/${limits.weeklyQueries} minute=${minuteUsage}/${limits.requestsPerMinute}`
+      );
+    }
+
+    return {
+      allowed,
+      limit: limits.weeklyQueries,
+      remaining: limits.weeklyQueries === -1 ? -1 : Math.max(0, limits.weeklyQueries - weeklyUsage),
+      resetAt: this.getWeeklyResetTime(),
+      planType,
+      limitType: !minuteAllowed ? "minute" : "weekly",
+      minuteLimit: limits.requestsPerMinute,
+      minuteRemaining: limits.requestsPerMinute === -1 ? -1 : Math.max(0, limits.requestsPerMinute - minuteUsage),
+      minuteResetAt: this.getMinuteResetTime(),
+    };
+  }
+
+  private async getPlanType(userId: string): Promise<string> {
+    return withD1Timeout(
+      async () => {
+        const result = await this.d1
+          .prepare(
+            `SELECT plan_type FROM user_subscriptions
+             WHERE user_id = ? AND status = 'active' LIMIT 1`
+          )
+          .bind(userId)
+          .first();
+        return (result?.plan_type as string) || "hobby";
+      },
+      "hobby",
+      "get_plan_type"
+    );
+  }
+
+  private async getUsageCount(
+    identifier: string,
+    period: "weekly" | "minute"
+  ): Promise<number> {
+    const since = period === "weekly"
+      ? this.getWeekStartTime().toISOString()
+      : new Date(Date.now() - 60_000).toISOString();
+
+    const operator = period === "weekly" ? ">=" : ">";
+
+    return withD1Timeout(
+      async () => {
+        const result = await this.d1
+          .prepare(
+            `SELECT
+              (SELECT COUNT(*) FROM search_logs WHERE user_id = ? AND created_at ${operator} ? AND status_code = 200) +
+              (SELECT COUNT(*) FROM fetch_logs WHERE user_id = ? AND created_at ${operator} ? AND status_code = 200) as total`
+          )
+          .bind(identifier, since, identifier, since)
+          .first();
+        return (result?.total as number) || 0;
+      },
+      0,
+      `get_${period}_usage`
+    );
+  }
+
+  private getWeekStartTime(): Date {
+    const now = new Date();
+    const start = new Date(now);
+    start.setDate(now.getDate() - now.getDay());
+    start.setHours(0, 0, 0, 0);
+    return start;
+  }
+
+  private getWeeklyResetTime(): string {
+    const now = new Date();
+    const next = new Date(now);
+    next.setDate(now.getDate() + (7 - now.getDay()));
+    next.setHours(0, 0, 0, 0);
+    return next.toISOString();
+  }
+
+  private getMinuteResetTime(): string {
+    const next = new Date();
+    next.setSeconds(0, 0);
+    next.setMinutes(next.getMinutes() + 1);
+    return next.toISOString();
+  }
+}

package/src/services/reranker.ts

@@ -0,0 +1,171 @@
+/**
+ * DeepInfra Reranker Service - MCP Optimized
+ * Dual-model fallback: 8B → 4B, 2 attempts each
+ */
+
+import { logger } from "../utils/logger.js";
+import { DEEPINFRA_CONFIG, DeepInfraService } from "./deepinfra-base.js";
+
+interface RerankerInput {
+  query: string;
+  documents: string[];
+  topN: number;
+}
+
+interface RerankerPayload {
+  queries: [string];
+  documents: string[];
+  top_n: number;
+}
+
+export interface RerankerResponse {
+  scores: number[];
+}
+
+export interface RankedDocument {
+  content: string;
+  originalIndex: number;
+  relevanceScore: number;
+}
+
+type ModelConfig = { endpoint: string; name: string };
+
+export class RerankerService extends DeepInfraService<
+  RerankerInput,
+  RerankerResponse,
+  RankedDocument[]
+> {
+  protected readonly endpoint =
+    `/v1/inference/${DEEPINFRA_CONFIG.RERANKER_MODEL_PRIMARY}`;
+
+  private static readonly MODELS: readonly ModelConfig[] = [
+    {
+      endpoint: `/v1/inference/${DEEPINFRA_CONFIG.RERANKER_MODEL_PRIMARY}`,
+      name: "8B",
+    },
+    {
+      endpoint: `/v1/inference/${DEEPINFRA_CONFIG.RERANKER_MODEL_FALLBACK}`,
+      name: "4B",
+    },
+  ];
+
+  private static readonly MAX_ATTEMPTS = 2;
+
+  async rerank(
+    query: string,
+    documents: string[],
+    topN: number
+  ): Promise<RankedDocument[]> {
+    if (!query?.trim()) throw new Error("Query cannot be empty for reranking");
+    if (!documents?.length)
+      throw new Error("Documents cannot be empty for reranking");
+
+    const validTopN = Math.min(topN, documents.length);
+    if (validTopN <= 0) throw new Error("top_n must be greater than 0");
+
+    return this.call(
+      { query: query.trim(), documents, topN: validTopN },
+      "Document reranking"
+    );
+  }
+
+  /**
+   * Dual-model fallback: 8B (2 attempts) → 4B (2 attempts) → fail
+   */
+  protected override async call(
+    input: RerankerInput,
+    operationName: string
+  ): Promise<RankedDocument[]> {
+    const startTime = Date.now();
+    const payload = this.buildPayload(input);
+    const errors: string[] = [];
+
+    for (const model of RerankerService.MODELS) {
+      const result = await this.tryModel(model, payload, input);
+      if (result.success) {
+        logger.info(
+          `${operationName} completed with ${model.name} (${this.elapsed(startTime)})`
+        );
+        return result.data!;
+      }
+      errors.push(`${model.name}: ${result.error!}`);
+      logger.warn(
+        `${model.name} model failed, ${model === RerankerService.MODELS[0] ? "switching to 4B" : "no more fallbacks"}`
+      );
+    }
+
+    logger.error(
+      `${operationName} failed after all attempts (${this.elapsed(startTime)}): ${errors.join(" | ")}`
+    );
+    throw new Error(`Reranking failed: ${errors.join(" | ")}`);
+  }
+
+  private async tryModel(
+    model: ModelConfig,
+    payload: RerankerPayload,
+    input: RerankerInput
+  ): Promise<{ success: boolean; data?: RankedDocument[]; error?: string }> {
+    let lastError = "";
+
+    for (let attempt = 1; attempt <= RerankerService.MAX_ATTEMPTS; attempt++) {
+      try {
+        const response = await this.singleRequest(model.endpoint, payload);
+        return { success: true, data: this.processResponse(response, input) };
+      } catch (e) {
+        lastError = e instanceof Error ? e.message : String(e);
+        logger.warn(
+          `${model.name} attempt ${attempt}/${RerankerService.MAX_ATTEMPTS} failed: ${lastError}`
+        );
+      }
+    }
+
+    return { success: false, error: lastError };
+  }
+
+  private elapsed(startTime: number): string {
+    return `${((Date.now() - startTime) / 1000).toFixed(1)}s`;
+  }
+
+  protected buildPayload(input: RerankerInput): RerankerPayload {
+    return {
+      queries: [input.query],
+      documents: input.documents,
+      top_n: input.topN,
+    };
+  }
+
+  protected processResponse(
+    response: RerankerResponse,
+    input: RerankerInput
+  ): RankedDocument[] {
+    if (!response.scores?.length) {
+      throw new Error("No reranking results received from DeepInfra API");
+    }
+
+    const count = Math.min(
+      input.topN,
+      response.scores.length,
+      input.documents.length
+    );
+
+    return input.documents
+      .map((content, index) => ({
+        content,
+        originalIndex: index,
+        relevanceScore: response.scores[index] ?? 0,
+      }))
+      .sort((a, b) => b.relevanceScore - a.relevanceScore)
+      .slice(0, count);
+  }
+
+  async healthCheck(): Promise<boolean> {
+    try {
+      return (await this.rerank("test query", ["test document"], 1)).length > 0;
+    } catch (error) {
+      logger.error(
+        `Reranker health check failed: ${error instanceof Error ? error.message : String(error)}`
+      );
+      return false;
+    }
+  }
+}

package/src/services/search-engine.ts

@@ -0,0 +1,333 @@
+/**
+ * Hybrid Search Engine for Apple Developer Documentation
+ *
+ * Advanced implementation combining Semantic Search for RAG with precise
+ * Keyword Search and Hybrid Search, optimized for developer documentation retrieval.
+ *
+ * Pipeline: Query → [Vector (4N) + Technical Term (4N)] → Merge → Title Merge → AI Rerank → Results
+ *
+ * Features:
+ * - 4N+4N hybrid candidate strategy
+ * - Semantic vector search with pgvector HNSW
+ * - Technical term search with PostgreSQL 'simple' configuration
+ * - Title-based content merging
+ * - AI reranking with Qwen3-Reranker-8B
+ */
+
+import type {
+  AdditionalUrl,
+  SearchOptions,
+  SearchResult,
+} from "../types/index.js";
+import { logger } from "../utils/logger.js";
+import type { DatabaseService } from "./database.js";
+import type { EmbeddingService } from "./embedding.js";
+import type { RerankerService } from "./reranker.js";
+
+export interface ParsedChunk {
+  content: string;
+  title: string | null;
+}
+
+export interface ProcessedResult {
+  id: string;
+  url: string;
+  title: string | null;
+  content: string;
+  contentLength: number;
+  chunk_index: number;
+  total_chunks: number;
+  mergedChunkIndices?: number[];
+}
+
+export interface RankedSearchResult {
+  id: string;
+  url: string;
+  title: string | null;
+  content: string;
+  chunk_index: number;
+  total_chunks: number;
+  mergedChunkIndices?: number[];
+  original_index: number;
+}
+
+export interface SearchEngineResult {
+  results: RankedSearchResult[];
+  additionalUrls: AdditionalUrl[];
+}
+
+export class SearchEngine {
+  constructor(
+    private database: DatabaseService,
+    private embedding: EmbeddingService,
+    private reranker: RerankerService
+  ) {}
+
+  /**
+   * Execute hybrid search optimized for Apple Developer Documentation
+   */
+  async search(
+    query: string,
+    options: SearchOptions = {}
+  ): Promise<SearchEngineResult> {
+    const { resultCount = 4 } = options;
+    return this.hybridSearchWithReranker(query, resultCount);
+  }
+
+  /**
+   * Hybrid search with 4N+4N candidate strategy
+   *
+   * 1. Parallel: Vector search (4N) + Technical term search (4N)
+   * 2. Merge and deduplicate by ID
+   * 3. Title-based content merging
+   * 4. AI reranking for optimal results
+   */
+  private async hybridSearchWithReranker(
+    query: string,
+    resultCount: number
+  ): Promise<SearchEngineResult> {
+    // Step 1: Parallel candidate retrieval (4N each, no minimum limit)
+    const candidateCount = resultCount * 4;
+
+    const [semanticResults, keywordResults] = await Promise.all([
+      this.getSemanticCandidates(query, candidateCount),
+      this.getKeywordCandidates(query, candidateCount),
+    ]);
+
+    // Step 2: Merge and deduplicate candidates
+    const mergedCandidates = this.mergeCandidates(
+      semanticResults,
+      keywordResults
+    );
+
+    // Step 3: Process results (title-based merging)
+    const processedResults = this.processResults(mergedCandidates);
+
+    // Step 4: AI reranking with fallback mechanism
+    let finalResults: RankedSearchResult[];
+
+    try {
+      const rankedDocuments = await this.reranker.rerank(
+        query,
+        processedResults.map((r) => r.content),
+        Math.min(resultCount, processedResults.length)
+      );
+
+      // Step 5: Map back to final results
+      finalResults = rankedDocuments.map((doc) => {
+        const processed = processedResults[doc.originalIndex];
+        return {
+          id: processed.id,
+          url: processed.url,
+          title: processed.title,
+          content: processed.content,
+          chunk_index: processed.chunk_index,
+          total_chunks: processed.total_chunks,
+          mergedChunkIndices: processed.mergedChunkIndices,
+          original_index: doc.originalIndex,
+        };
+      });
+    } catch (error) {
+      logger.error(
+        `Reranking failed, falling back to original order (query_length: ${query.length}, candidates: ${processedResults.length}): ${error instanceof Error ? error.message : String(error)}`
+      );
+
+      // Fallback: use original order, truncate to requested count
+      finalResults = processedResults
+        .slice(0, resultCount)
+        .map((processed, index) => ({
+          id: processed.id,
+          url: processed.url,
+          title: processed.title,
+          content: processed.content,
+          chunk_index: processed.chunk_index,
+          total_chunks: processed.total_chunks,
+          mergedChunkIndices: processed.mergedChunkIndices,
+          original_index: index,
+        }));
+
+      logger.warn(
+        `Reranking failed, using original order with ${finalResults.length} results`
+      );
+    }
+
+    // Collect additional URLs
+    const additionalUrls = this.collectAdditionalUrls(
+      processedResults,
+      finalResults
+    );
+
+    return { results: finalResults, additionalUrls };
+  }
+
+  /**
+   * Retrieve semantic search candidates with error handling
+   */
+  private async getSemanticCandidates(
+    query: string,
+    resultCount: number
+  ): Promise<SearchResult[]> {
+    const startTime = Date.now();
+
+    try {
+      const queryEmbedding = await this.embedding.createEmbedding(query);
+      const results = await this.database.semanticSearch(queryEmbedding, {
+        resultCount,
+      });
+
+      logger.info(
+        `Semantic search completed (${((Date.now() - startTime) / 1000).toFixed(1)}s): ${results.length} results`
+      );
+
+      return results;
+    } catch (error) {
+      logger.error(
+        `Semantic search failed, falling back to keyword-only: ${error instanceof Error ? error.message : String(error)}`
+      );
+      return [];
+    }
+  }
+
+  /**
+   * Retrieve keyword search candidates with error handling
+   */
+  private async getKeywordCandidates(
+    query: string,
+    resultCount: number
+  ): Promise<SearchResult[]> {
+    const startTime = Date.now();
+
+    try {
+      const results = await this.database.keywordSearch(query, {
+        resultCount,
+      });
+
+      logger.info(
+        `Keyword search completed (${((Date.now() - startTime) / 1000).toFixed(1)}s): ${results.length} results`
+      );
+
+      return results;
+    } catch (error) {
+      logger.error(
+        `Keyword search failed: ${error instanceof Error ? error.message : String(error)}`
+      );
+      return [];
+    }
+  }
+
+  /**
+   * Merge and deduplicate candidates from semantic and keyword search
+   */
+  private mergeCandidates(
+    semanticResults: SearchResult[],
+    keywordResults: SearchResult[]
+  ): SearchResult[] {
+    const seen = new Set<string>();
+
+    // Prioritize semantic results, then add unique keyword results
+    return [
+      ...semanticResults.filter((result) => {
+        if (seen.has(result.id)) return false;
+        seen.add(result.id);
+        return true;
+      }),
+      ...keywordResults.filter((result) => {
+        if (seen.has(result.id)) return false;
+        seen.add(result.id);
+        return true;
+      }),
+    ];
+  }
+
+  /**
+   * Collect additional URLs from processed results
+   */
+  private collectAdditionalUrls(
+    processedResults: ProcessedResult[],
+    finalResults: RankedSearchResult[]
+  ): AdditionalUrl[] {
+    const finalUrls = new Set(finalResults.map((r) => r.url));
+
+    return processedResults
+      .filter((r) => !finalUrls.has(r.url))
+      .reduce((urls, r) => {
+        if (!urls.some((u) => u.url === r.url)) {
+          urls.push({
+            url: r.url,
+            title: r.title,
+            characterCount: r.contentLength,
+          });
+        }
+        return urls;
+      }, [] as AdditionalUrl[])
+      .slice(0, 10);
+  }
+
+  /**
+   * Process RAG candidates through title-based merging
+   */
+  private processResults(candidates: SearchResult[]): ProcessedResult[] {
+    // Step 1: Merge by title
+    return this.mergeByTitle(candidates);
+  }
+
+  private parseChunk(content: string, title: string | null): ParsedChunk {
+    // Since data migration is complete, content is now plain text
+    // and title comes from the dedicated title field
+    return {
+      title: title || "",
+      content: content,
+    };
+  }
+
+  private mergeByTitle(results: SearchResult[]): ProcessedResult[] {
+    const titleGroups = new Map<string, SearchResult[]>();
+
+    // Group by title
+    for (const result of results) {
+      const { title } = this.parseChunk(result.content, result.title);
+      const titleKey = title || "untitled";
+      if (!titleGroups.has(titleKey)) {
+        titleGroups.set(titleKey, []);
+      }
+      titleGroups.get(titleKey)!.push(result);
+    }
+
+    return Array.from(titleGroups.entries()).map(([title, group]) => {
+      const primary = group[0];
+
+      // Sort and merge chunks by original index to maintain proper content order
+      const chunkIndices = group
+        .map((r) => r.chunk_index)
+        .sort((a, b) => a - b);
+      const mergedContent = group
+        .sort((a, b) => a.chunk_index - b.chunk_index)
+        .map((r) => this.parseChunk(r.content, r.title).content)
+        .join("\n\n---\n\n");
+
+      // Detect complete document merging
+      const isCompleteDocument =
+        chunkIndices.length === primary.total_chunks &&
+        chunkIndices.every((idx, i) => idx === i);
+
+      // Determine final chunk representation
+      const [chunk_index, total_chunks] =
+        chunkIndices.length === 1
+          ? [chunkIndices[0], primary.total_chunks]
+          : isCompleteDocument
+            ? [0, 1]
+            : [Math.min(...chunkIndices), primary.total_chunks];
+
+      return {
+        id: primary.id,
+        url: primary.url,
+        title,
+        content: mergedContent,
+        mergedChunkIndices: chunkIndices.length > 1 ? chunkIndices : undefined,
+        contentLength: mergedContent.length,
+        chunk_index,
+        total_chunks,
+      };
+    });
+  }
+}