File size: 2,384 Bytes
a5784e9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
/**
 * Model Capabilities Hook
 * Fetches model capabilities from backend single source of truth
 */

import { useQuery } from '@tanstack/react-query';

// Types matching backend response
export type ThinkingType = 'level' | 'budget' | 'none';

export interface CategoryCapabilities {
  thinkingType: ThinkingType;
  levels?: string[];
  defaultLevel?: string;
  alwaysOn?: boolean;
  budgetRange?: [number, number];
  defaultBudget?: number;
  supportsGoogleSearch?: boolean;
}

export interface ModelCapabilitiesResponse {
  categories: Record<string, CategoryCapabilities>;
  matchers: Array<{ pattern: string; category: string }>;
}

export type ModelCategory = 
  | 'gemini3Flash'
  | 'gemini3Pro'
  | 'gemini25Pro'
  | 'gemini25Flash'
  | 'gemini2'
  | 'other';

async function fetchCapabilities(): Promise<ModelCapabilitiesResponse> {
  const response = await fetch('/api/model-capabilities');
  if (!response.ok) {
    throw new Error('Failed to fetch model capabilities');
  }
  return response.json();
}

/**
 * Hook to fetch and cache model capabilities
 * Returns capabilities and helper functions
 */
export function useModelCapabilities() {
  const { data, isLoading, error } = useQuery({
    queryKey: ['model-capabilities'],
    queryFn: fetchCapabilities,
    staleTime: Infinity, // Capabilities rarely change
    gcTime: Infinity,
  });

  /**
   * Get category for a model ID using backend-provided matchers
   */
  const getModelCategory = (modelId: string): ModelCategory => {
    if (!data) return 'other';
    
    const modelLower = modelId.toLowerCase();
    
    for (const matcher of data.matchers) {
      const regex = new RegExp(matcher.pattern, 'i');
      if (regex.test(modelLower)) {
        return matcher.category as ModelCategory;
      }
    }
    
    return 'other';
  };

  /**
   * Get capabilities for a specific category
   */
  const getCategoryCapabilities = (category: ModelCategory): CategoryCapabilities | undefined => {
    return data?.categories[category];
  };

  /**
   * Get capabilities for a model ID
   */
  const getModelCapabilities = (modelId: string): CategoryCapabilities | undefined => {
    const category = getModelCategory(modelId);
    return getCategoryCapabilities(category);
  };

  return {
    data,
    isLoading,
    error,
    getModelCategory,
    getCategoryCapabilities,
    getModelCapabilities,
  };
}