diff --git a/api/chat_model.py b/api/chat_model.py index f7cbf2fc4..1d3c17b05 100644 --- a/api/chat_model.py +++ b/api/chat_model.py @@ -7,6 +7,7 @@ class ChatMessage(BaseModel): role: str # 'user' or 'assistant' content: str + mode: Literal["normal", "deep_research"] = Field(default="normal") class ChatCompletionRequest(BaseModel): """ @@ -26,6 +27,13 @@ class ChatCompletionRequest(BaseModel): model: Optional[str] = Field(None, description="Model name for the specified provider") language: Optional[str] = Field("en", description="Language for content generation (e.g., 'en', 'ja', 'zh', 'es', 'kr', 'vi')") + + research_iteration: int = Field( + default=1, + ge=1, + description="Current deep research iteration (1-based). Only used when the request is in deep_research mode.", + ) + excluded_dirs: List[str] = Field( default_factory=list, description="List or newline-separated string of directories to exclude from processing", diff --git a/api/simple_chat.py b/api/simple_chat.py index 7db123ec1..9c4ae413a 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -116,22 +116,11 @@ async def chat_completions_stream(request: ChatCompletionRequest): ) # Check if this is a Deep Research request - is_deep_research = False - research_iteration = 1 - - # Process messages to detect Deep Research requests - for msg in request.messages: - if hasattr(msg, 'content') and msg.content and "[DEEP RESEARCH]" in msg.content: - is_deep_research = True - # Only remove the tag from the last message - if msg == request.messages[-1]: - # Remove the Deep Research tag - msg.content = msg.content.replace("[DEEP RESEARCH]", "").strip() + is_deep_research = last_message.mode == "deep_research" # Count research iterations if this is a Deep Research request if is_deep_research: - research_iteration = sum(1 for msg in request.messages if msg.role == 'assistant') + 1 - logger.info(f"Deep Research request detected - iteration {research_iteration}") + logger.info("Deep Research request detected - iteration %d",request.research_iteration) # Check if this is a continuation request if "continue" in last_message.content.lower() and "research" in last_message.content.lower(): @@ -139,7 +128,7 @@ async def chat_completions_stream(request: ChatCompletionRequest): original_topic = None for msg in request.messages: if msg.role == "user" and "continue" not in msg.content.lower(): - original_topic = msg.content.replace("[DEEP RESEARCH]", "").strip() + original_topic = msg.content.strip() logger.info(f"Found original research topic: {original_topic}") break @@ -219,10 +208,10 @@ async def chat_completions_stream(request: ChatCompletionRequest): # Create system prompt if is_deep_research: # Check if this is the first iteration - is_first_iteration = research_iteration == 1 + is_first_iteration = request.research_iteration == 1 # Check if this is the final iteration - is_final_iteration = research_iteration >= 5 + is_final_iteration = request.research_iteration >= 5 if is_first_iteration: system_prompt = DEEP_RESEARCH_FIRST_ITERATION_PROMPT.format( @@ -243,7 +232,7 @@ async def chat_completions_stream(request: ChatCompletionRequest): repo_type=repo_type, repo_url=repo_url, repo_name=repo_name, - research_iteration=research_iteration, + research_iteration=request.research_iteration, language_name=language_name ) else: diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index 32e37f7d9..d7ddfda2c 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -119,22 +119,11 @@ async def handle_websocket_chat(websocket: WebSocket): ) # Check if this is a Deep Research request - is_deep_research = False - research_iteration = 1 - - # Process messages to detect Deep Research requests - for msg in request.messages: - if hasattr(msg, 'content') and msg.content and "[DEEP RESEARCH]" in msg.content: - is_deep_research = True - # Only remove the tag from the last message - if msg == request.messages[-1]: - # Remove the Deep Research tag - msg.content = msg.content.replace("[DEEP RESEARCH]", "").strip() + is_deep_research = last_message.mode == "deep_research" # Count research iterations if this is a Deep Research request if is_deep_research: - research_iteration = sum(1 for msg in request.messages if msg.role == 'assistant') + 1 - logger.info(f"Deep Research request detected - iteration {research_iteration}") + logger.info("Deep Research request detected - iteration %d", request.research_iteration) # Check if this is a continuation request if "continue" in last_message.content.lower() and "research" in last_message.content.lower(): @@ -142,7 +131,7 @@ async def handle_websocket_chat(websocket: WebSocket): original_topic = None for msg in request.messages: if msg.role == "user" and "continue" not in msg.content.lower(): - original_topic = msg.content.replace("[DEEP RESEARCH]", "").strip() + original_topic = msg.content.strip() logger.info(f"Found original research topic: {original_topic}") break @@ -222,10 +211,10 @@ async def handle_websocket_chat(websocket: WebSocket): # Create system prompt if is_deep_research: # Check if this is the first iteration - is_first_iteration = research_iteration == 1 + is_first_iteration = request.research_iteration == 1 # Check if this is the final iteration - is_final_iteration = research_iteration >= 5 + is_final_iteration = request.research_iteration >= 5 if is_first_iteration: system_prompt = DEEP_RESEARCH_FIRST_ITERATION_PROMPT.format( @@ -246,7 +235,7 @@ async def handle_websocket_chat(websocket: WebSocket): repo_type=repo_type, repo_url=repo_url, repo_name=repo_name, - research_iteration=research_iteration, + research_iteration=request.research_iteration, language_name=language_name ) else: diff --git a/src/components/Ask.tsx b/src/components/Ask.tsx index 056afcd23..e91f0b066 100644 --- a/src/components/Ask.tsx +++ b/src/components/Ask.tsx @@ -24,6 +24,7 @@ interface Provider { interface Message { role: 'user' | 'assistant' | 'system'; content: string; + mode?: 'normal' | 'deep_research'; } interface ResearchStage { @@ -299,7 +300,8 @@ const Ask: React.FC = ({ }, { role: 'user', - content: '[DEEP RESEARCH] Continue the research' + content: 'Continue the research', + mode: 'deep_research' } ]; @@ -317,10 +319,15 @@ const Ask: React.FC = ({ const requestBody: ChatCompletionRequest = { repo_url: getRepoUrl(repoInfo), type: repoInfo.type, - messages: newHistory.map(msg => ({ role: msg.role as 'user' | 'assistant', content: msg.content })), + messages: newHistory.map(msg => ({ + role: msg.role as 'user' | 'assistant', + content: msg.content, + mode: msg.mode ?? 'normal', + })), provider: selectedProvider, model: isCustomSelectedModel ? customSelectedModel : selectedModel, - language: language + language: language, + research_iteration: newIteration }; // Add tokens if available @@ -548,7 +555,8 @@ const Ask: React.FC = ({ // Create initial message const initialMessage: Message = { role: 'user', - content: deepResearch ? `[DEEP RESEARCH] ${question}` : question + content: question, + mode: deepResearch ? 'deep_research' : 'normal' }; // Set initial conversation history @@ -559,10 +567,15 @@ const Ask: React.FC = ({ const requestBody: ChatCompletionRequest = { repo_url: getRepoUrl(repoInfo), type: repoInfo.type, - messages: newHistory.map(msg => ({ role: msg.role as 'user' | 'assistant', content: msg.content })), + messages: newHistory.map(msg => ({ + role: msg.role as 'user' | 'assistant', + content: msg.content, + mode: msg.mode ?? 'normal' + })), provider: selectedProvider, model: isCustomSelectedModel ? customSelectedModel : selectedModel, - language: language + language: language, + research_iteration: deepResearch ? 1 : undefined }; // Add tokens if available diff --git a/src/utils/websocketClient.ts b/src/utils/websocketClient.ts index b29c1b104..f3b22a8a1 100644 --- a/src/utils/websocketClient.ts +++ b/src/utils/websocketClient.ts @@ -17,6 +17,7 @@ const getWebSocketUrl = () => { export interface ChatMessage { role: 'user' | 'assistant' | 'system'; content: string; + mode?: 'normal' | 'deep_research'; } export interface ChatCompletionRequest { @@ -28,6 +29,7 @@ export interface ChatCompletionRequest { provider?: string; model?: string; language?: string; + research_iteration?: number; excluded_dirs?: string; excluded_files?: string; }