From cf26ffd47acebdaa53c0209d73b6ea2506917ad9 Mon Sep 17 00:00:00 2001
From: suyao <sy20010504@gmail.com>
Date: Thu, 22 May 2025 22:51:22 +0800
Subject: [PATCH] feat: add support for 'grok' provider in web search
 functionality

- Enhanced `isWebSearchModel` to recognize 'grok' as a valid web search model.
- Updated `getOpenAIWebSearchParams` to return specific search parameters for 'grok'.
- Modified `OpenAIProvider` to handle citations from 'grok' in web search results.
- Added 'grok' to the `WebSearchSource` enum for citation formatting.
---
 src/renderer/src/config/models.ts                  | 14 ++++++++++++++
 .../src/providers/AiProvider/OpenAIProvider.ts     | 12 ++++++++++++
 src/renderer/src/store/messageBlock.ts             |  1 +
 src/renderer/src/types/index.ts                    |  3 ++-
 4 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/renderer/src/config/models.ts b/src/renderer/src/config/models.ts
index 2ade2c1a74..5f82d39d70 100644
--- a/src/renderer/src/config/models.ts
+++ b/src/renderer/src/config/models.ts
@@ -2570,6 +2570,10 @@ export function isWebSearchModel(model: Model): boolean {
     return true
   }
 
+  if (provider.id === 'grok') {
+    return true
+  }
+
   return false
 }
 
@@ -2600,6 +2604,16 @@ export function getOpenAIWebSearchParams(assistant: Assistant, model: Model): Re
     if (assistant.enableWebSearch) {
       const webSearchTools = getWebSearchTools(model)
 
+      if (model.provider === 'grok') {
+        return {
+          search_parameters: {
+            mode: 'auto',
+            return_citations: true,
+            sources: [{ type: 'web' }, { type: 'x' }, { type: 'news' }]
+          }
+        }
+      }
+
       if (model.provider === 'hunyuan') {
         return { enable_enhancement: true, citation: true, search_info: true }
       }
diff --git a/src/renderer/src/providers/AiProvider/OpenAIProvider.ts b/src/renderer/src/providers/AiProvider/OpenAIProvider.ts
index 8d27569ee4..3e3a6e432a 100644
--- a/src/renderer/src/providers/AiProvider/OpenAIProvider.ts
+++ b/src/renderer/src/providers/AiProvider/OpenAIProvider.ts
@@ -777,6 +777,18 @@ export default class OpenAIProvider extends BaseOpenAIProvider {
                   }
                 } as LLMWebSearchCompleteChunk)
               }
+              if (assistant.model?.provider === 'grok') {
+                const citations = originalFinishRawChunk.citations
+                if (citations) {
+                  onChunk({
+                    type: ChunkType.LLM_WEB_SEARCH_COMPLETE,
+                    llm_web_search: {
+                      results: citations,
+                      source: WebSearchSource.GROK
+                    }
+                  } as LLMWebSearchCompleteChunk)
+                }
+              }
               if (assistant.model?.provider === 'perplexity') {
                 const citations = originalFinishRawChunk.citations
                 if (citations) {
diff --git a/src/renderer/src/store/messageBlock.ts b/src/renderer/src/store/messageBlock.ts
index c9cc55cec3..cf9515fe00 100644
--- a/src/renderer/src/store/messageBlock.ts
+++ b/src/renderer/src/store/messageBlock.ts
@@ -160,6 +160,7 @@ const formatCitationsFromBlock = (block: CitationMessageBlock | undefined): Cita
             }
           }) || []
         break
+      case WebSearchSource.GROK:
       case WebSearchSource.OPENROUTER:
       case WebSearchSource.PERPLEXITY:
         formattedCitations =
diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts
index a269d32676..921b846a3f 100644
--- a/src/renderer/src/types/index.ts
+++ b/src/renderer/src/types/index.ts
@@ -473,7 +473,8 @@ export enum WebSearchSource {
   PERPLEXITY = 'perplexity',
   QWEN = 'qwen',
   HUNYUAN = 'hunyuan',
-  ZHIPU = 'zhipu'
+  ZHIPU = 'zhipu',
+  GROK = 'grok'
 }
 
 export type WebSearchResponse = {