@lobehub/chat 1.90.3 → 1.91.0

package/src/server/services/search/impls/jina/type.ts

@@ -0,0 +1,26 @@
+export interface JinaSearchParameters {
+  q: string;
+}
+
+interface JinaUsage {
+  tokens: number;
+}
+
+interface JinaMeta {
+  usage: JinaUsage;
+}
+
+interface JinaData {
+  content?: string;
+  description?: string;
+  title: string;
+  url: string;
+  usage?: JinaUsage;
+}
+
+export interface JinaResponse {
+  code?: number;
+  data: JinaData[];
+  meta?: JinaMeta;
+  status?: number;
+}

package/src/server/services/search/impls/tavily/index.ts

@@ -0,0 +1,124 @@
+import { TRPCError } from '@trpc/server';
+import debug from 'debug';
+import urlJoin from 'url-join';
+
+import { SearchParams, UniformSearchResponse, UniformSearchResult } from '@/types/tool/search';
+
+import { SearchServiceImpl } from '../type';
+import { TavilySearchParameters, TavilyResponse } from './type';
+
+const log = debug('lobe-search:Tavily');
+
+/**
+ * Tavily implementation of the search service
+ * Primarily used for web crawling
+ */
+export class TavilyImpl implements SearchServiceImpl {
+  private get apiKey(): string | undefined {
+    return process.env.TAVILY_API_KEY;
+  }
+
+  private get baseUrl(): string {
+    // Assuming the base URL is consistent with the crawl endpoint
+    return 'https://api.tavily.com';
+  }
+
+  async query(query: string, params: SearchParams = {}): Promise<UniformSearchResponse> {
+    log('Starting Tavily query with query: "%s", params: %o', query, params);
+    const endpoint = urlJoin(this.baseUrl, '/search');
+
+    const defaultQueryParams: TavilySearchParameters = {
+      include_answer: false,
+      include_image_descriptions: true,
+      include_images: false,
+      include_raw_content: false,
+      max_results: 15,
+      query,
+      search_depth: process.env.TAVILY_SEARCH_DEPTH || 'basic' // basic or advanced
+    };
+
+    let body: TavilySearchParameters = {
+      ...defaultQueryParams,
+      time_range:
+        params?.searchTimeRange && params.searchTimeRange !== 'anytime'
+          ? params.searchTimeRange
+          : undefined,
+      topic:
+        // Tavily 只支持 news 和 general 两种类型
+        params?.searchCategories?.filter(cat => ['news', 'general'].includes(cat))?.[0],
+    };
+
+    log('Constructed request body: %o', body);
+
+    let response: Response;
+    const startAt = Date.now();
+    let costTime = 0;
+    try {
+      log('Sending request to endpoint: %s', endpoint);
+      response = await fetch(endpoint, {
+        body: JSON.stringify(body),
+        headers: {
+          'Authorization': this.apiKey ? `Bearer ${this.apiKey}` : '',
+          'Content-Type': 'application/json',
+        },
+        method: 'POST',
+      });
+      log('Received response with status: %d', response.status);
+      costTime = Date.now() - startAt;
+    } catch (error) {
+      log.extend('error')('Tavily fetch error: %o', error);
+      throw new TRPCError({
+        cause: error,
+        code: 'SERVICE_UNAVAILABLE',
+        message: 'Failed to connect to Tavily.',
+      });
+    }
+
+    if (!response.ok) {
+      const errorBody = await response.text();
+      log.extend('error')(
+        `Tavily request failed with status ${response.status}: %s`,
+        errorBody.length > 200 ? `${errorBody.slice(0, 200)}...` : errorBody,
+      );
+      throw new TRPCError({
+        cause: errorBody,
+        code: 'SERVICE_UNAVAILABLE',
+        message: `Tavily request failed: ${response.statusText}`,
+      });
+    }
+
+    try {
+      const tavilyResponse = (await response.json()) as TavilyResponse;
+
+      log('Parsed Tavily response: %o', tavilyResponse);
+
+      const mappedResults = (tavilyResponse.results || []).map(
+        (result): UniformSearchResult => ({
+          category: body.topic || 'general', // Default category
+          content: result.content || '', // Prioritize content, fallback to snippet
+          engines: ['tavily'], // Use 'tavily' as the engine name
+          parsedUrl: result.url ? new URL(result.url).hostname : '', // Basic URL parsing
+          score: result.score || 0, // Default score to 0 if undefined
+          title: result.title || '',
+          url: result.url,
+        }),
+      );
+
+      log('Mapped %d results to SearchResult format', mappedResults.length);
+
+      return {
+        costTime,
+        query: query,
+        resultNumbers: mappedResults.length,
+        results: mappedResults,
+      };
+    } catch (error) {
+      log.extend('error')('Error parsing Tavily response: %o', error);
+      throw new TRPCError({
+        cause: error,
+        code: 'INTERNAL_SERVER_ERROR',
+        message: 'Failed to parse Tavily response.',
+      });
+    }
+  }
+}

package/src/server/services/search/impls/tavily/type.ts

@@ -0,0 +1,36 @@
+export interface TavilySearchParameters {
+  chunks_per_source?: number;
+  days?: number;
+  exclude_domains?: string[];
+  include_answer?: boolean | string;
+  include_domains?: string[];
+  include_image_descriptions?: boolean;
+  include_images?: boolean;
+  include_raw_content?: boolean;
+  max_results?: number;
+  query: string;
+  search_depth?: string;
+  time_range?: string;
+  topic?: string;
+}
+
+interface TavilyImages {
+  description?: string;
+  url: string;
+}
+
+interface TavilyResults {
+  content?: string;
+  raw_content?: string | null;
+  score?: number;
+  title?: string;
+  url: string;
+}
+
+export interface TavilyResponse {
+  answer?: string;
+  images?: TavilyImages[];
+  query: string;
+  response_time: number;
+  results: TavilyResults[];
+}

package/src/store/chat/slices/aiChat/actions/__tests__/generateAIChat.test.ts

@@ -671,14 +671,16 @@ describe('chatMessage actions', () => {
       (chatService.createAssistantMessage as Mock).mockResolvedValue(aiResponse);
       const spy = vi.spyOn(chatService, 'createAssistantMessageStream');
       // 模拟消息创建
-      (messageService.createMessage as Mock).mockResolvedValue('assistant-message-id');
+      const createMessageSpyOn = vi
+        .spyOn(messageService, 'createMessage')
+        .mockResolvedValue('assistant-message-id');
 
       await act(async () => {
         await result.current.internal_coreProcessMessage(messages, userMessage.id);
       });
 
       // 验证是否创建了代表 AI 响应的消息
-      expect(messageService.createMessage).toHaveBeenCalledWith(
+      expect(createMessageSpyOn).toHaveBeenCalledWith(
         expect.objectContaining({
           role: 'assistant',
           content: LOADING_FLAT,

package/src/store/chat/slices/message/action.test.ts

@@ -369,11 +369,12 @@ describe('chatMessage actions', () => {
       const messageId = 'message-id';
       const newContent = 'Updated content';
 
+      const spy = vi.spyOn(messageService, 'updateMessage');
       await act(async () => {
         await result.current.internal_updateMessageContent(messageId, newContent);
       });
 
-      expect(messageService.updateMessage).toHaveBeenCalledWith(messageId, { content: newContent });
+      expect(spy).toHaveBeenCalledWith(messageId, { content: newContent });
     });
 
     it('should dispatch message update action', async () => {

package/src/store/chat/slices/topic/action.test.ts

@@ -264,7 +264,8 @@ describe('topic action', () => {
       const topicId = 'topic-id';
       const newTitle = 'Updated Topic Title';
       // Mock the topicService.updateTitle to resolve immediately
-      (topicService.updateTopic as Mock).mockResolvedValue(undefined);
+
+      const spyOn = vi.spyOn(topicService, 'updateTopic');
 
       const { result } = renderHook(() => useChatStore());
 
@@ -276,7 +277,7 @@ describe('topic action', () => {
       });
 
       // Verify that the topicService.updateTitle was called with correct parameters
-      expect(topicService.updateTopic).toHaveBeenCalledWith(topicId, {
+      expect(spyOn).toHaveBeenCalledWith(topicId, {
         title: 'Updated Topic Title',
       });