diff --git a/api/chat_model.py b/api/chat_model.py new file mode 100644 index 000000000..f7cbf2fc4 --- /dev/null +++ b/api/chat_model.py @@ -0,0 +1,57 @@ +from typing import Optional, List, Literal +from urllib.parse import unquote + +from pydantic import BaseModel, Field, field_validator + +# Models for the API +class ChatMessage(BaseModel): + role: str # 'user' or 'assistant' + content: str + +class ChatCompletionRequest(BaseModel): + """ + Model for requesting a chat completion. + """ + repo_url: str = Field(..., description="URL of the repository to query") + messages: List[ChatMessage] = Field(..., description="List of chat messages") + filePath: Optional[str] = Field(None, description="Optional path to a file in the repository to include in the prompt") + token: Optional[str] = Field(None, description="Personal access token for private repositories") + type: Optional[Literal["github", "gitlab", "bitbucket"]] = Field( + "github", + description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')", + ) + + # model parameters + provider: str = Field("google", description="Model provider (google, openai, openrouter, ollama, bedrock, azure, dashscope)") + 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')") + excluded_dirs: List[str] = Field( + default_factory=list, + description="List or newline-separated string of directories to exclude from processing", + ) + excluded_files: List[str] = Field( + default_factory=list, + description="List or newline-separated string of file patterns to exclude from processing", + ) + included_dirs: List[str] = Field( + default_factory=list, + description="List or newline-separated string of directories to include exclusively", + ) + included_files: List[str] = Field( + default_factory=list, + description="List or newline-separated string of file patterns to include exclusively", + ) + + @field_validator( + "excluded_dirs", + "excluded_files", + "included_dirs", + "included_files", + mode="before", + ) + @classmethod + def validate_path(cls, value: list[str] | str) -> list[str]: + if isinstance(value, str): + value = [unquote(path) for path in value.strip().split("\n") if path] + return value diff --git a/api/simple_chat.py b/api/simple_chat.py index 9076f71ab..a1abe5786 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -1,12 +1,10 @@ import logging -from typing import List, Optional, Callable -from urllib.parse import unquote +from typing import Callable from functools import partial from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse -from pydantic import BaseModel, Field from api.chat import ChatStreamer, prompt_builder, is_token_limit_error from api.config import get_model_config, configs @@ -18,6 +16,7 @@ DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT, SIMPLE_CHAT_SYSTEM_PROMPT ) +from api.chat_model import ChatCompletionRequest # Configure logging from api.logging_config import setup_logging @@ -41,31 +40,6 @@ allow_headers=["*"], # Allows all headers ) -# Models for the API -class ChatMessage(BaseModel): - role: str # 'user' or 'assistant' - content: str - -class ChatCompletionRequest(BaseModel): - """ - Model for requesting a chat completion. - """ - repo_url: str = Field(..., description="URL of the repository to query") - messages: List[ChatMessage] = Field(..., description="List of chat messages") - filePath: Optional[str] = Field(None, description="Optional path to a file in the repository to include in the prompt") - token: Optional[str] = Field(None, description="Personal access token for private repositories") - type: Optional[str] = Field("github", description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')") - - # model parameters - provider: str = Field("google", description="Model provider (google, openai, openrouter, ollama, bedrock, azure, dashscope)") - 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')") - excluded_dirs: Optional[str] = Field(None, description="Comma-separated list of directories to exclude from processing") - excluded_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to exclude from processing") - included_dirs: Optional[str] = Field(None, description="Comma-separated list of directories to include exclusively") - included_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to include exclusively") - @app.post("/chat/completions/stream") async def chat_completions_stream(request: ChatCompletionRequest): @@ -87,25 +61,24 @@ async def chat_completions_stream(request: ChatCompletionRequest): request_rag = RAG(provider=request.provider, model=request.model) # Extract custom file filter parameters if provided - excluded_dirs = None - excluded_files = None - included_dirs = None - included_files = None - if request.excluded_dirs: - excluded_dirs = [unquote(dir_path) for dir_path in request.excluded_dirs.split('\n') if dir_path.strip()] - logger.info(f"Using custom excluded directories: {excluded_dirs}") + logger.info(f"Using custom excluded directories: {request.excluded_dirs}") if request.excluded_files: - excluded_files = [unquote(file_pattern) for file_pattern in request.excluded_files.split('\n') if file_pattern.strip()] - logger.info(f"Using custom excluded files: {excluded_files}") + logger.info(f"Using custom excluded files: {request.excluded_files}") if request.included_dirs: - included_dirs = [unquote(dir_path) for dir_path in request.included_dirs.split('\n') if dir_path.strip()] - logger.info(f"Using custom included directories: {included_dirs}") + logger.info(f"Using custom included directories: {request.included_dirs}") if request.included_files: - included_files = [unquote(file_pattern) for file_pattern in request.included_files.split('\n') if file_pattern.strip()] - logger.info(f"Using custom included files: {included_files}") - - request_rag.prepare_retriever(request.repo_url, request.type, request.token, excluded_dirs, excluded_files, included_dirs, included_files) + logger.info(f"Using custom included files: {request.included_files}") + + request_rag.prepare_retriever( + request.repo_url, + request.type, + request.token, + excluded_dirs=request.excluded_dirs, + excluded_files=request.excluded_files, + included_dirs=request.included_dirs, + included_files=request.included_files, + ) logger.info(f"Retriever prepared for {request.repo_url}") except ValueError as e: if "No valid documents with embeddings found" in str(e): diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index 2d7e45b71..22e8e9ca5 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -1,11 +1,8 @@ import logging from collections.abc import AsyncIterator, Callable -from typing import List, Optional -from urllib.parse import unquote from functools import partial from fastapi import WebSocket, WebSocketDisconnect -from pydantic import BaseModel, Field from api.chat import ChatStreamer, prompt_builder, is_token_limit_error from api.config import ( @@ -20,6 +17,7 @@ DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT, SIMPLE_CHAT_SYSTEM_PROMPT, ) +from api.chat_model import ChatCompletionRequest # Configure logging from api.logging_config import setup_logging @@ -28,34 +26,6 @@ logger = logging.getLogger(__name__) -# Models for the API -class ChatMessage(BaseModel): - role: str # 'user' or 'assistant' - content: str - -class ChatCompletionRequest(BaseModel): - """ - Model for requesting a chat completion. - """ - repo_url: str = Field(..., description="URL of the repository to query") - messages: List[ChatMessage] = Field(..., description="List of chat messages") - filePath: Optional[str] = Field(None, description="Optional path to a file in the repository to include in the prompt") - token: Optional[str] = Field(None, description="Personal access token for private repositories") - type: Optional[str] = Field("github", description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')") - - # model parameters - provider: str = Field( - "google", - description="Model provider (google, openai, openrouter, ollama, bedrock, azure, dashscope)", - ) - 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')") - excluded_dirs: Optional[str] = Field(None, description="Comma-separated list of directories to exclude from processing") - excluded_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to exclude from processing") - included_dirs: Optional[str] = Field(None, description="Comma-separated list of directories to include exclusively") - included_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to include exclusively") - async def handle_websocket_chat(websocket: WebSocket): """ Handle WebSocket connection for chat completions. @@ -84,25 +54,24 @@ async def handle_websocket_chat(websocket: WebSocket): request_rag = RAG(provider=request.provider, model=request.model) # Extract custom file filter parameters if provided - excluded_dirs = None - excluded_files = None - included_dirs = None - included_files = None - if request.excluded_dirs: - excluded_dirs = [unquote(dir_path) for dir_path in request.excluded_dirs.split('\n') if dir_path.strip()] - logger.info(f"Using custom excluded directories: {excluded_dirs}") + logger.info(f"Using custom excluded directories: {request.excluded_dirs}") if request.excluded_files: - excluded_files = [unquote(file_pattern) for file_pattern in request.excluded_files.split('\n') if file_pattern.strip()] - logger.info(f"Using custom excluded files: {excluded_files}") + logger.info(f"Using custom excluded files: {request.excluded_files}") if request.included_dirs: - included_dirs = [unquote(dir_path) for dir_path in request.included_dirs.split('\n') if dir_path.strip()] - logger.info(f"Using custom included directories: {included_dirs}") + logger.info(f"Using custom included directories: {request.included_dirs}") if request.included_files: - included_files = [unquote(file_pattern) for file_pattern in request.included_files.split('\n') if file_pattern.strip()] - logger.info(f"Using custom included files: {included_files}") - - request_rag.prepare_retriever(request.repo_url, request.type, request.token, excluded_dirs, excluded_files, included_dirs, included_files) + logger.info(f"Using custom included files: {request.included_files}") + + request_rag.prepare_retriever( + request.repo_url, + request.type, + request.token, + excluded_dirs=request.excluded_dirs, + excluded_files=request.excluded_files, + included_dirs=request.included_dirs, + included_files=request.included_files, + ) logger.info(f"Retriever prepared for {request.repo_url}") except ValueError as e: if "No valid documents with embeddings found" in str(e):