From d61eb44106822511462ef7e11ac79fed59b2c31d Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 26 Jul 2026 17:36:15 +0800 Subject: [PATCH 1/7] refactor RAG, move rag related data_pipeline, ollama model validation, and RAG to dedicated rag module --- api/ollama_patch.py | 49 --- api/rag.py | 449 ---------------------- api/rag/__init__.py | 7 + api/{data_pipeline.py => rag/pipeline.py} | 282 +------------- api/rag/rag.py | 406 +++++++++++++++++++ api/repository.py | 306 +++++++++++++++ api/simple_chat.py | 24 +- api/websocket_wiki.py | 22 +- 8 files changed, 754 insertions(+), 791 deletions(-) delete mode 100644 api/ollama_patch.py delete mode 100644 api/rag.py create mode 100644 api/rag/__init__.py rename api/{data_pipeline.py => rag/pipeline.py} (70%) create mode 100644 api/rag/rag.py create mode 100644 api/repository.py diff --git a/api/ollama_patch.py b/api/ollama_patch.py deleted file mode 100644 index cc04b7317..000000000 --- a/api/ollama_patch.py +++ /dev/null @@ -1,49 +0,0 @@ -import logging - -import os - -# Configure logging -from api.logging_config import setup_logging - -setup_logging() -logger = logging.getLogger(__name__) - -class OllamaModelNotFoundError(Exception): - """Custom exception for when Ollama model is not found""" - pass - -def check_ollama_model_exists(model_name: str, ollama_host: str = None) -> bool: - """ - Check if an Ollama model exists before attempting to use it. - - Args: - model_name: Name of the model to check - ollama_host: Ollama host URL, defaults to localhost:11434 - - Returns: - bool: True if model exists, False otherwise - """ - import ollama - import httpx - if ollama_host is None: - ollama_host = os.getenv("OLLAMA_HOST", "http://localhost:11434") - try: - # Remove /api prefix if present and add it back - if ollama_host.endswith('/api'): - ollama_host = ollama_host[:-4] - ret: ollama.ListResponse = ollama.Client(host=ollama_host, timeout=5).list() - is_available = any(model_name == model.model for model in ret.models) - if is_available: - logger.info("Ollama model '%s' is available", model_name) - else: - logger.warning( - "Ollama model '%s' is not available. Available models: %s. ", - model_name, - str([model.model for model in ret.models])) - return is_available - except (httpx.ConnectTimeout, ConnectionError) as e: - logger.warning(f"Could not connect to Ollama to check models: {e}") - return False - except Exception as e: - logger.warning(f"Error checking Ollama model availability: {e}") - return False diff --git a/api/rag.py b/api/rag.py deleted file mode 100644 index 157ed3460..000000000 --- a/api/rag.py +++ /dev/null @@ -1,449 +0,0 @@ -import logging -import asyncio -import os -from dataclasses import dataclass -from typing import Any, List, Tuple, Dict -from uuid import uuid4 - -import adalflow as adal - -from api.tools.embedder import get_embedder -from api.prompts import RAG_SYSTEM_PROMPT as system_prompt, RAG_TEMPLATE - -# Create our own implementation of the conversation classes -@dataclass -class UserQuery: - query_str: str - -@dataclass -class AssistantResponse: - response_str: str - -@dataclass -class DialogTurn: - id: str - user_query: UserQuery - assistant_response: AssistantResponse - -class CustomConversation: - """Custom implementation of Conversation to fix the list assignment index out of range error""" - - def __init__(self): - self.dialog_turns = [] - - def append_dialog_turn(self, dialog_turn): - """Safely append a dialog turn to the conversation""" - if not hasattr(self, 'dialog_turns'): - self.dialog_turns = [] - self.dialog_turns.append(dialog_turn) - -# Import other adalflow components -from adalflow.components.retriever.faiss_retriever import FAISSRetriever -from api.config import configs -from api.data_pipeline import DatabaseManager - -# Configure logging -logger = logging.getLogger(__name__) - -# Maximum token limit for embedding models -MAX_INPUT_TOKENS = 7500 # Safe threshold below 8192 token limit - -# Maximum concurrent RAG preparing count -_RAG_PREPARE_SEMAPHORE: asyncio.Semaphore | None = None - - -def _get_rag_semaphore() -> asyncio.Semaphore: - global _RAG_PREPARE_SEMAPHORE - if _RAG_PREPARE_SEMAPHORE is None: - _RAG_PREPARE_SEMAPHORE = asyncio.Semaphore( - int(os.environ.get("DEEPWIKI_MAX_CONCURRENT_RAG", "4")) - ) - assert isinstance(_RAG_PREPARE_SEMAPHORE, asyncio.Semaphore) - return _RAG_PREPARE_SEMAPHORE - - -class Memory(adal.core.component.DataComponent): - """Simple conversation management with a list of dialog turns.""" - - def __init__(self): - super().__init__() - # Use our custom implementation instead of the original Conversation class - self.current_conversation = CustomConversation() - - def call(self) -> Dict: - """Return the conversation history as a dictionary.""" - all_dialog_turns = {} - try: - # Check if dialog_turns exists and is a list - if hasattr(self.current_conversation, 'dialog_turns'): - if self.current_conversation.dialog_turns: - logger.info(f"Memory content: {len(self.current_conversation.dialog_turns)} turns") - for i, turn in enumerate(self.current_conversation.dialog_turns): - if hasattr(turn, 'id') and turn.id is not None: - all_dialog_turns[turn.id] = turn - logger.info(f"Added turn {i+1} with ID {turn.id} to memory") - else: - logger.warning(f"Skipping invalid turn object in memory: {turn}") - else: - logger.info("Dialog turns list exists but is empty") - else: - logger.info("No dialog_turns attribute in current_conversation") - # Try to initialize it - self.current_conversation.dialog_turns = [] - except Exception as e: - logger.error(f"Error accessing dialog turns: {str(e)}") - # Try to recover - try: - self.current_conversation = CustomConversation() - logger.info("Recovered by creating new conversation") - except Exception as e2: - logger.error(f"Failed to recover: {str(e2)}") - - logger.info(f"Returning {len(all_dialog_turns)} dialog turns from memory") - return all_dialog_turns - - def add_dialog_turn(self, user_query: str, assistant_response: str) -> bool: - """ - Add a dialog turn to the conversation history. - - Args: - user_query: The user's query - assistant_response: The assistant's response - - Returns: - bool: True if successful, False otherwise - """ - try: - # Create a new dialog turn using our custom implementation - dialog_turn = DialogTurn( - id=str(uuid4()), - user_query=UserQuery(query_str=user_query), - assistant_response=AssistantResponse(response_str=assistant_response), - ) - - # Make sure the current_conversation has the append_dialog_turn method - if not hasattr(self.current_conversation, 'append_dialog_turn'): - logger.warning("current_conversation does not have append_dialog_turn method, creating new one") - # Initialize a new conversation if needed - self.current_conversation = CustomConversation() - - # Ensure dialog_turns exists - if not hasattr(self.current_conversation, 'dialog_turns'): - logger.warning("dialog_turns not found, initializing empty list") - self.current_conversation.dialog_turns = [] - - # Safely append the dialog turn - self.current_conversation.dialog_turns.append(dialog_turn) - logger.info(f"Successfully added dialog turn, now have {len(self.current_conversation.dialog_turns)} turns") - return True - - except Exception as e: - logger.error(f"Error adding dialog turn: {str(e)}") - # Try to recover by creating a new conversation - try: - self.current_conversation = CustomConversation() - dialog_turn = DialogTurn( - id=str(uuid4()), - user_query=UserQuery(query_str=user_query), - assistant_response=AssistantResponse(response_str=assistant_response), - ) - self.current_conversation.dialog_turns.append(dialog_turn) - logger.info("Recovered from error by creating new conversation") - return True - except Exception as e2: - logger.error(f"Failed to recover from error: {str(e2)}") - return False - - -from dataclasses import dataclass, field - -@dataclass -class RAGAnswer(adal.DataClass): - rationale: str = field(default="", metadata={"desc": "Chain of thoughts for the answer."}) - answer: str = field(default="", metadata={"desc": "Answer to the user query, formatted in markdown for beautiful rendering with react-markdown. DO NOT include ``` triple backticks fences at the beginning or end of your answer."}) - - __output_fields__ = ["rationale", "answer"] - -class RAG(adal.Component): - """RAG with one repo. - If you want to load a new repos, call prepare_retriever(repo_url_or_path) first.""" - - def __init__(self, provider="google", model=None, use_s3: bool = False): # noqa: F841 - use_s3 is kept for compatibility - """ - Initialize the RAG component. - - Args: - provider: Model provider to use (google, openai, openrouter, ollama) - model: Model name to use with the provider - use_s3: Whether to use S3 for database storage (default: False) - """ - super().__init__() - - self.provider = provider - self.model = model - - # Import the helper functions - from api.config import get_embedder_config, get_embedder_type - - # Determine embedder type based on current configuration - self.embedder_type = get_embedder_type() - self.is_ollama_embedder = (self.embedder_type == 'ollama') # Backward compatibility - - # Check if Ollama model exists before proceeding - if self.is_ollama_embedder: - from api.ollama_patch import check_ollama_model_exists - from api.config import get_embedder_config - - embedder_config = get_embedder_config() - if embedder_config and embedder_config.get("model_kwargs", {}).get("model"): - model_name = embedder_config["model_kwargs"]["model"] - if not check_ollama_model_exists(model_name): - raise Exception(f"Ollama model '{model_name}' not found. Please run 'ollama pull {model_name}' to install it.") - - # Initialize components - self.memory = Memory() - self.embedder = get_embedder(embedder_type=self.embedder_type) - self.initialize_db_manager() - - def initialize_db_manager(self): - """Initialize the database manager with local storage""" - self.db_manager = DatabaseManager() - self.transformed_docs = [] - - def _validate_and_filter_embeddings(self, documents: List) -> List: - """ - Validate embeddings and filter out documents with invalid or mismatched embedding sizes. - - Args: - documents: List of documents with embeddings - - Returns: - List of documents with valid embeddings of consistent size - """ - if not documents: - logger.warning("No documents provided for embedding validation") - return [] - - valid_documents = [] - embedding_sizes = {} - - # First pass: collect all embedding sizes and count occurrences - for i, doc in enumerate(documents): - if not hasattr(doc, 'vector') or doc.vector is None: - logger.warning(f"Document {i} has no embedding vector, skipping") - continue - - try: - if isinstance(doc.vector, list): - embedding_size = len(doc.vector) - elif hasattr(doc.vector, 'shape'): - embedding_size = doc.vector.shape[0] if len(doc.vector.shape) == 1 else doc.vector.shape[-1] - elif hasattr(doc.vector, '__len__'): - embedding_size = len(doc.vector) - else: - logger.warning(f"Document {i} has invalid embedding vector type: {type(doc.vector)}, skipping") - continue - - if embedding_size == 0: - logger.warning(f"Document {i} has empty embedding vector, skipping") - continue - - embedding_sizes[embedding_size] = embedding_sizes.get(embedding_size, 0) + 1 - - except Exception as e: - logger.warning(f"Error checking embedding size for document {i}: {str(e)}, skipping") - continue - - if not embedding_sizes: - logger.error("No valid embeddings found in any documents") - return [] - - # Find the most common embedding size (this should be the correct one) - target_size = max(embedding_sizes.keys(), key=lambda k: embedding_sizes[k]) - logger.info(f"Target embedding size: {target_size} (found in {embedding_sizes[target_size]} documents)") - - # Log all embedding sizes found - for size, count in embedding_sizes.items(): - if size != target_size: - logger.warning(f"Found {count} documents with incorrect embedding size {size}, will be filtered out") - - # Second pass: filter documents with the target embedding size - for i, doc in enumerate(documents): - if not hasattr(doc, 'vector') or doc.vector is None: - continue - - try: - if isinstance(doc.vector, list): - embedding_size = len(doc.vector) - elif hasattr(doc.vector, 'shape'): - embedding_size = doc.vector.shape[0] if len(doc.vector.shape) == 1 else doc.vector.shape[-1] - elif hasattr(doc.vector, '__len__'): - embedding_size = len(doc.vector) - else: - continue - - if embedding_size == target_size: - valid_documents.append(doc) - else: - # Log which document is being filtered out - file_path = getattr(doc, 'meta_data', {}).get('file_path', f'document_{i}') - logger.warning(f"Filtering out document '{file_path}' due to embedding size mismatch: {embedding_size} != {target_size}") - - except Exception as e: - file_path = getattr(doc, 'meta_data', {}).get('file_path', f'document_{i}') - logger.warning(f"Error validating embedding for document '{file_path}': {str(e)}, skipping") - continue - - logger.info(f"Embedding validation complete: {len(valid_documents)}/{len(documents)} documents have valid embeddings") - - if len(valid_documents) == 0: - logger.error("No documents with valid embeddings remain after filtering") - elif len(valid_documents) < len(documents): - filtered_count = len(documents) - len(valid_documents) - logger.warning(f"Filtered out {filtered_count} documents due to embedding issues") - - return valid_documents - - def prepare_retriever(self, repo_url_or_path: str, type: str = "github", access_token: str = None, - excluded_dirs: List[str] = None, excluded_files: List[str] = None, - included_dirs: List[str] = None, included_files: List[str] = None): - """ - Prepare the retriever for a repository. - Will load database from local storage if available. - - Args: - repo_url_or_path: URL or local path to the repository - access_token: Optional access token for private repositories - excluded_dirs: Optional list of directories to exclude from processing - excluded_files: Optional list of file patterns to exclude from processing - included_dirs: Optional list of directories to include exclusively - included_files: Optional list of file patterns to include exclusively - """ - self.initialize_db_manager() - self.repo_url_or_path = repo_url_or_path - self.transformed_docs = self.db_manager.prepare_database( - repo_url_or_path, - type, - access_token, - embedder_type=self.embedder_type, - excluded_dirs=excluded_dirs, - excluded_files=excluded_files, - included_dirs=included_dirs, - included_files=included_files - ) - logger.info(f"Loaded {len(self.transformed_docs)} documents for retrieval") - - # Validate and filter embeddings to ensure consistent sizes - self.transformed_docs = self._validate_and_filter_embeddings(self.transformed_docs) - - if not self.transformed_docs: - raise ValueError("No valid documents with embeddings found. Cannot create retriever.") - - logger.info(f"Using {len(self.transformed_docs)} documents with valid embeddings for retrieval") - - try: - # Use the appropriate embedder for retrieval - self.retriever = FAISSRetriever( - **configs["retriever"], - embedder=self.embedder, - documents=self.transformed_docs, - document_map_func=lambda doc: doc.vector, - ) - logger.info("FAISS retriever created successfully") - except Exception as e: - logger.error(f"Error creating FAISS retriever: {str(e)}") - # Try to provide more specific error information - if "All embeddings should be of the same size" in str(e): - logger.error("Embedding size validation failed. This suggests there are still inconsistent embedding sizes.") - # Log embedding sizes for debugging - sizes = [] - for i, doc in enumerate(self.transformed_docs[:10]): # Check first 10 docs - if hasattr(doc, 'vector') and doc.vector is not None: - try: - if isinstance(doc.vector, list): - size = len(doc.vector) - elif hasattr(doc.vector, 'shape'): - size = doc.vector.shape[0] if len(doc.vector.shape) == 1 else doc.vector.shape[-1] - elif hasattr(doc.vector, '__len__'): - size = len(doc.vector) - else: - size = "unknown" - sizes.append(f"doc_{i}: {size}") - except Exception: - sizes.append(f"doc_{i}: error") - logger.error(f"Sample embedding sizes: {', '.join(sizes)}") - raise - - async def aprepare_retriever( - self, - repo_url_or_path: str, - type: str = "github", - access_token: str | None = None, - excluded_dirs: list[str] | None = None, - excluded_files: list[str] | None = None, - included_dirs: list[str] | None = None, - included_files: list[str] | None = None, - ): - """Async version of the original `prepare_retriever`. - - Reuse the synchronous `prepare_retriever` implementation, but runs it in - a worker thread via `asyncio.to_thread` so that blocking operations (such - as git.clone, file io, embedding calls) do not stall the outer event loop. - Concurrency is bounded by a module-level semaphore, set by system variable - 'DEEPWIKI_MAX_CONCURRENT_RAG'. - - Args: - repo_url_or_path: URL or local path to the repository - access_token: Optional access token for private repositories - excluded_dirs: Optional list of directories to exclude from processing - excluded_files: Optional list of file patterns to exclude from processing - included_dirs: Optional list of directories to include exclusively - included_files: Optional list of file patterns to include exclusively - """ - async with _get_rag_semaphore(): - return await asyncio.to_thread( - self.prepare_retriever, - repo_url_or_path, - type=type, - access_token=access_token, - excluded_dirs=excluded_dirs, - excluded_files=excluded_files, - included_dirs=included_dirs, - included_files=included_files, - ) - - def call(self, query: str, language: str = "en") -> Tuple[List]: - """ - Process a query using RAG. - - Args: - query: The user's query - - Returns: - Tuple of (RAGAnswer, retrieved_documents) - """ - try: - retrieved_documents = self.retriever(query) - - # Fill in the documents - retrieved_documents[0].documents = [ - self.transformed_docs[doc_index] - for doc_index in retrieved_documents[0].doc_indices - ] - - return retrieved_documents - - except Exception as e: - logger.error(f"Error in RAG call: {str(e)}") - - # Create error response - error_response = RAGAnswer( - rationale="Error occurred while processing the query.", - answer=f"I apologize, but I encountered an error while processing your question. Please try again or rephrase your question." - ) - return error_response, [] - - async def acall(self, query: str, language: str = "en") -> tuple[list]: - """Async version of the original `call` method. - """ - return await asyncio.to_thread(self.call, query, language) diff --git a/api/rag/__init__.py b/api/rag/__init__.py new file mode 100644 index 000000000..4df56ed48 --- /dev/null +++ b/api/rag/__init__.py @@ -0,0 +1,7 @@ +from api.rag.pipeline import count_tokens +from api.rag.rag import RAG + +__all__ = [ + "RAG", + "count_tokens", +] diff --git a/api/data_pipeline.py b/api/rag/pipeline.py similarity index 70% rename from api/data_pipeline.py rename to api/rag/pipeline.py index 16d15a84d..4a1c1b1dd 100644 --- a/api/data_pipeline.py +++ b/api/rag/pipeline.py @@ -1,21 +1,17 @@ -import adalflow as adal -from adalflow.core.types import Document, List -from adalflow.components.data_process import TextSplitter, ToEmbeddings +import logging import os import subprocess -import json -import tiktoken -import logging -import base64 -import glob from pathlib import Path -from adalflow.utils import get_adalflow_default_root_path +from urllib.parse import quote, urlparse, urlunparse + +import adalflow as adal +import tiktoken +from adalflow.components.data_process import TextSplitter, ToEmbeddings from adalflow.core.db import LocalDB -from api.config import configs, DEFAULT_EXCLUDED_DIRS, DEFAULT_EXCLUDED_FILES -from urllib.parse import urlparse, urlunparse, quote -import requests -from requests.exceptions import RequestException +from adalflow.core.types import Document, List +from adalflow.utils import get_adalflow_default_root_path +from api.config import DEFAULT_EXCLUDED_DIRS, DEFAULT_EXCLUDED_FILES, configs from api.tools.embedder import get_embedder # Configure logging @@ -426,266 +422,6 @@ def transform_documents_and_save_to_db( db.save_state(filepath=db_path) return db -def get_github_file_content(repo_url: str, file_path: str, access_token: str = None) -> str: - """ - Retrieves the content of a file from a GitHub repository using the GitHub API. - Supports both public GitHub (github.com) and GitHub Enterprise (custom domains). - - Args: - repo_url (str): The URL of the GitHub repository - (e.g., "https://github.com/username/repo" or "https://github.company.com/username/repo") - file_path (str): The path to the file within the repository (e.g., "src/main.py") - access_token (str, optional): GitHub personal access token for private repositories - - Returns: - str: The content of the file as a string - - Raises: - ValueError: If the file cannot be fetched or if the URL is not a valid GitHub URL - """ - try: - # Parse the repository URL to support both github.com and enterprise GitHub - parsed_url = urlparse(repo_url) - if not parsed_url.scheme or not parsed_url.netloc: - raise ValueError("Not a valid GitHub repository URL") - - # Check if it's a GitHub-like URL structure - path_parts = parsed_url.path.strip('/').split('/') - if len(path_parts) < 2: - raise ValueError("Invalid GitHub URL format - expected format: https://domain/owner/repo") - - owner = path_parts[-2] - repo = path_parts[-1].replace(".git", "") - - # Determine the API base URL - if parsed_url.netloc == "github.com": - # Public GitHub - api_base = "https://api.github.com" - else: - # GitHub Enterprise - API is typically at https://domain/api/v3/ - api_base = f"{parsed_url.scheme}://{parsed_url.netloc}/api/v3" - - # Use GitHub API to get file content - # The API endpoint for getting file content is: /repos/{owner}/{repo}/contents/{path} - api_url = f"{api_base}/repos/{owner}/{repo}/contents/{file_path}" - - # Fetch file content from GitHub API - headers = {} - if access_token: - headers["Authorization"] = f"token {access_token}" - logger.info(f"Fetching file content from GitHub API: {api_url}") - try: - response = requests.get(api_url, headers=headers) - response.raise_for_status() - except RequestException as e: - raise ValueError(f"Error fetching file content: {e}") - try: - content_data = response.json() - except json.JSONDecodeError: - raise ValueError("Invalid response from GitHub API") - - # Check if we got an error response - if "message" in content_data and "documentation_url" in content_data: - raise ValueError(f"GitHub API error: {content_data['message']}") - - # GitHub API returns file content as base64 encoded string - if "content" in content_data and "encoding" in content_data: - if content_data["encoding"] == "base64": - # The content might be split into lines, so join them first - content_base64 = content_data["content"].replace("\n", "") - content = base64.b64decode(content_base64).decode("utf-8") - return content - else: - raise ValueError(f"Unexpected encoding: {content_data['encoding']}") - else: - raise ValueError("File content not found in GitHub API response") - - except Exception as e: - raise ValueError(f"Failed to get file content: {str(e)}") - -def get_gitlab_file_content(repo_url: str, file_path: str, access_token: str = None) -> str: - """ - Retrieves the content of a file from a GitLab repository (cloud or self-hosted). - - Args: - repo_url (str): The GitLab repo URL (e.g., "https://gitlab.com/username/repo" or "http://localhost/group/project") - file_path (str): File path within the repository (e.g., "src/main.py") - access_token (str, optional): GitLab personal access token - - Returns: - str: File content - - Raises: - ValueError: If anything fails - """ - try: - # Parse and validate the URL - parsed_url = urlparse(repo_url) - if not parsed_url.scheme or not parsed_url.netloc: - raise ValueError("Not a valid GitLab repository URL") - - gitlab_domain = f"{parsed_url.scheme}://{parsed_url.netloc}" - if parsed_url.port not in (None, 80, 443): - gitlab_domain += f":{parsed_url.port}" - path_parts = parsed_url.path.strip("/").split("/") - if len(path_parts) < 2: - raise ValueError("Invalid GitLab URL format — expected something like https://gitlab.domain.com/group/project") - - # Build project path and encode for API - project_path = "/".join(path_parts).replace(".git", "") - encoded_project_path = quote(project_path, safe='') - - # Encode file path - encoded_file_path = quote(file_path, safe='') - - # Try to get the default branch from the project info - default_branch = None - try: - project_info_url = f"{gitlab_domain}/api/v4/projects/{encoded_project_path}" - project_headers = {} - if access_token: - project_headers["PRIVATE-TOKEN"] = access_token - - project_response = requests.get(project_info_url, headers=project_headers) - if project_response.status_code == 200: - project_data = project_response.json() - default_branch = project_data.get('default_branch', 'main') - logger.info(f"Found default branch: {default_branch}") - else: - logger.warning(f"Could not fetch project info, using 'main' as default branch") - default_branch = 'main' - except Exception as e: - logger.warning(f"Error fetching project info: {e}, using 'main' as default branch") - default_branch = 'main' - - api_url = f"{gitlab_domain}/api/v4/projects/{encoded_project_path}/repository/files/{encoded_file_path}/raw?ref={default_branch}" - # Fetch file content from GitLab API - headers = {} - if access_token: - headers["PRIVATE-TOKEN"] = access_token - logger.info(f"Fetching file content from GitLab API: {api_url}") - try: - response = requests.get(api_url, headers=headers) - response.raise_for_status() - content = response.text - except RequestException as e: - raise ValueError(f"Error fetching file content: {e}") - - # Check for GitLab error response (JSON instead of raw file) - if content.startswith("{") and '"message":' in content: - try: - error_data = json.loads(content) - if "message" in error_data: - raise ValueError(f"GitLab API error: {error_data['message']}") - except json.JSONDecodeError: - pass - - return content - - except Exception as e: - raise ValueError(f"Failed to get file content: {str(e)}") - -def get_bitbucket_file_content(repo_url: str, file_path: str, access_token: str = None) -> str: - """ - Retrieves the content of a file from a Bitbucket repository using the Bitbucket API. - - Args: - repo_url (str): The URL of the Bitbucket repository (e.g., "https://bitbucket.org/username/repo") - file_path (str): The path to the file within the repository (e.g., "src/main.py") - access_token (str, optional): Bitbucket personal access token for private repositories - - Returns: - str: The content of the file as a string - """ - try: - # Extract owner and repo name from Bitbucket URL - if not (repo_url.startswith("https://bitbucket.org/") or repo_url.startswith("http://bitbucket.org/")): - raise ValueError("Not a valid Bitbucket repository URL") - - parts = repo_url.rstrip('/').split('/') - if len(parts) < 5: - raise ValueError("Invalid Bitbucket URL format") - - owner = parts[-2] - repo = parts[-1].replace(".git", "") - - # Try to get the default branch from the repository info - default_branch = None - try: - repo_info_url = f"https://api.bitbucket.org/2.0/repositories/{owner}/{repo}" - repo_headers = {} - if access_token: - repo_headers["Authorization"] = f"Bearer {access_token}" - - repo_response = requests.get(repo_info_url, headers=repo_headers) - if repo_response.status_code == 200: - repo_data = repo_response.json() - default_branch = repo_data.get('mainbranch', {}).get('name', 'main') - logger.info(f"Found default branch: {default_branch}") - else: - logger.warning(f"Could not fetch repository info, using 'main' as default branch") - default_branch = 'main' - except Exception as e: - logger.warning(f"Error fetching repository info: {e}, using 'main' as default branch") - default_branch = 'main' - - # Use Bitbucket API to get file content - # The API endpoint for getting file content is: /2.0/repositories/{owner}/{repo}/src/{branch}/{path} - api_url = f"https://api.bitbucket.org/2.0/repositories/{owner}/{repo}/src/{default_branch}/{file_path}" - - # Fetch file content from Bitbucket API - headers = {} - if access_token: - headers["Authorization"] = f"Bearer {access_token}" - logger.info(f"Fetching file content from Bitbucket API: {api_url}") - try: - response = requests.get(api_url, headers=headers) - if response.status_code == 200: - content = response.text - elif response.status_code == 404: - raise ValueError("File not found on Bitbucket. Please check the file path and repository.") - elif response.status_code == 401: - raise ValueError("Unauthorized access to Bitbucket. Please check your access token.") - elif response.status_code == 403: - raise ValueError("Forbidden access to Bitbucket. You might not have permission to access this file.") - elif response.status_code == 500: - raise ValueError("Internal server error on Bitbucket. Please try again later.") - else: - response.raise_for_status() - content = response.text - return content - except RequestException as e: - raise ValueError(f"Error fetching file content: {e}") - - except Exception as e: - raise ValueError(f"Failed to get file content: {str(e)}") - - -def get_file_content(repo_url: str, file_path: str, repo_type: str = None, access_token: str = None) -> str: - """ - Retrieves the content of a file from a Git repository (GitHub or GitLab). - - Args: - repo_type (str): Type of repository - repo_url (str): The URL of the repository - file_path (str): The path to the file within the repository - access_token (str, optional): Access token for private repositories - - Returns: - str: The content of the file as a string - - Raises: - ValueError: If the file cannot be fetched or if the URL is not valid - """ - if repo_type == "github": - return get_github_file_content(repo_url, file_path, access_token) - elif repo_type == "gitlab": - return get_gitlab_file_content(repo_url, file_path, access_token) - elif repo_type == "bitbucket": - return get_bitbucket_file_content(repo_url, file_path, access_token) - else: - raise ValueError("Unsupported repository type. Only GitHub, GitLab, and Bitbucket are supported.") - class DatabaseManager: """ Manages the creation, loading, transformation, and persistence of LocalDB instances. diff --git a/api/rag/rag.py b/api/rag/rag.py new file mode 100644 index 000000000..069a84c49 --- /dev/null +++ b/api/rag/rag.py @@ -0,0 +1,406 @@ +import asyncio +import logging +import os +from collections import defaultdict +from collections.abc import Sized +from dataclasses import dataclass, field +from uuid import uuid4 + +import adalflow as adal +from adalflow.components.retriever.faiss_retriever import FAISSRetriever +from adalflow.core.types import AssistantResponse, DialogTurn, Document, UserQuery + +from api.config import configs +from api.rag.pipeline import DatabaseManager +from api.tools.embedder import get_embedder + +# Configure logging +logger = logging.getLogger(__name__) + +# Maximum token limit for embedding models +MAX_INPUT_TOKENS = 7500 # Safe threshold below 8192 token limit + +# Maximum concurrent RAG preparing count +_RAG_PREPARE_SEMAPHORE: asyncio.Semaphore | None = None + + +def _get_rag_semaphore() -> asyncio.Semaphore: + global _RAG_PREPARE_SEMAPHORE + if _RAG_PREPARE_SEMAPHORE is None: + _RAG_PREPARE_SEMAPHORE = asyncio.Semaphore( + int(os.environ.get("DEEPWIKI_MAX_CONCURRENT_RAG", "4")) + ) + assert isinstance(_RAG_PREPARE_SEMAPHORE, asyncio.Semaphore) + return _RAG_PREPARE_SEMAPHORE + + +def check_ollama_model_exists(model_name: str, ollama_host: str | None = None) -> bool: + """ + Check if an Ollama model exists before attempting to use it. + + Args: + model_name: Name of the model to check + ollama_host: Ollama host URL, defaults to localhost:11434 + + Returns: + bool: True if model exists, False otherwise + """ + import httpx + import ollama + + if ollama_host is None: + ollama_host = os.getenv("OLLAMA_HOST", "http://localhost:11434") + try: + # Remove /api prefix if present and add it back + ollama_host = ollama_host.removesuffix("/api") + ret: ollama.ListResponse = ollama.Client(host=ollama_host, timeout=5).list() + is_available = any(model_name == model.model for model in ret.models) + if is_available: + logger.info("Ollama model '%s' is available", model_name) + else: + logger.warning( + "Ollama model '%s' is not available. Available models: %s. ", + model_name, + str([model.model for model in ret.models]), + ) + return is_available + except (httpx.ConnectTimeout, ConnectionError) as e: + logger.warning(f"Could not connect to Ollama to check models: {e}") + return False + except Exception as e: + logger.warning(f"Error checking Ollama model availability: {e}") + return False + + +class CustomConversation(list[DialogTurn]): + """Custom implementation of Conversation to fix the list assignment index out of range error""" + + +class Memory(adal.core.component.DataComponent): + """Simple conversation management with a list of dialog turns.""" + + def __init__(self): + super().__init__() + # Use our custom implementation instead of the original Conversation class + self.current_conversation = CustomConversation() + + def call(self) -> dict: + """Return the conversation history as a dictionary.""" + all_dialog_turns = ( + {} + if not self.current_conversation + else { + dialog_turn.id: dialog_turn for dialog_turn in self.current_conversation + } + ) + logger.info(f"Returning {len(all_dialog_turns)} dialog turns from memory") + return all_dialog_turns + + def add_dialog_turn(self, user_query: str, assistant_response: str) -> None: + """ + Add a dialog turn to the conversation history. + + Args: + user_query: The user's query + assistant_response: The assistant's response + + """ + # Create a new dialog turn using our custom implementation + dialog_turn = DialogTurn( + id=str(uuid4()), + user_query=UserQuery(query_str=user_query), + assistant_response=AssistantResponse(response_str=assistant_response), + ) + + # Safely append the dialog turn + self.current_conversation.append(dialog_turn) + logger.info( + "Successfully added dialog turn, now have %d turns", + len(self.current_conversation), + ) + + +@dataclass +class RAGAnswer(adal.DataClass): + rationale: str = field( + default="", metadata={"desc": "Chain of thoughts for the answer."} + ) + answer: str = field( + default="", + metadata={ + "desc": "Answer to the user query, formatted in markdown for beautiful rendering with react-markdown. DO NOT include ``` triple backticks fences at the beginning or end of your answer." + }, + ) + + __output_fields__ = ["rationale", "answer"] + + +def _get_document_vector_size(document: Document) -> int | None: + if hasattr(document.vector, "shape"): + embedding_size = ( + document.vector.shape[0] + if len(document.vector.shape) == 1 + else document.vector.shape[-1] + ) + elif isinstance(document.vector, Sized): + embedding_size = len(document.vector) + else: + embedding_size = None + + return embedding_size + + +class RAG(adal.Component): + """RAG with one repo. + If you want to load a new repos, call prepare_retriever(repo_url_or_path) first.""" + + def __init__(self, provider="google", model=None, use_s3: bool = False): # noqa: F841 - use_s3 is kept for compatibility + """ + Initialize the RAG component. + + Args: + provider: Model provider to use (google, openai, openrouter, ollama) + model: Model name to use with the provider + use_s3: Whether to use S3 for database storage (default: False) + """ + super().__init__() + + self.provider = provider + self.model = model + + # Import the helper functions + from api.config import get_embedder_type + + # Determine embedder type based on current configuration + self.embedder_type = get_embedder_type() + self.is_ollama_embedder = ( + self.embedder_type == "ollama" + ) # Backward compatibility + + # Check if Ollama model exists before proceeding + if self.is_ollama_embedder: + from api.config import get_embedder_config + + embedder_config = get_embedder_config() + if embedder_config and embedder_config.get("model_kwargs", {}).get("model"): + model_name = embedder_config["model_kwargs"]["model"] + if not check_ollama_model_exists(model_name): + raise ValueError( + f"Ollama model '{model_name}' not found. Please run 'ollama pull {model_name}' to install it." + ) + + # Initialize components + self.memory = Memory() + self.embedder = get_embedder(embedder_type=self.embedder_type) + self.initialize_db_manager() + + def initialize_db_manager(self): + """Initialize the database manager with local storage""" + self.db_manager = DatabaseManager() + self.transformed_docs = [] + + @staticmethod + def _validate_and_filter_embeddings(documents: list[Document]) -> list: + """ + Validate embeddings and filter out documents with invalid or mismatched embedding sizes. + + Args: + documents: List of documents with embeddings + + Returns: + List of documents with valid embeddings of consistent size + """ + if not documents: + logger.warning("No documents provided for embedding validation") + return [] + + docs_embeddings = defaultdict(list) + + for doc, embed_size in filter( + lambda x: isinstance(x[0], Document) and bool(x[1]), + ((x, _get_document_vector_size(x)) for x in documents), + ): + docs_embeddings[embed_size].append(doc) + if not docs_embeddings: + logger.error("No valid embeddings found in any documents") + return [] + + target_size = max(docs_embeddings, key=lambda x: len(docs_embeddings[x])) + logger.info( + "Target embedding size: %s (found in %s documents)", + target_size, + len(docs_embeddings[target_size]), + ) + + valid_documents = docs_embeddings.pop(target_size) + + if docs_embeddings: + for embed_size, docs_list in docs_embeddings.items(): + logger.warning( + "Found %s documents with incorrect embedding size %s, will be filtered out.", + len(docs_list), + str(embed_size), + ) + + if not valid_documents: + logger.warning("No documents with valid embeddings remained after filtering") + else: + logger.info( + "Embedding validation complete: %d/%d documents have valid embeddings.", + len(valid_documents), + len(documents), + ) + return valid_documents + + def prepare_retriever( + self, + repo_url_or_path: str, + type: str = "github", + access_token: str | None = None, + excluded_dirs: list[str] | None = None, + excluded_files: list[str] | None = None, + included_dirs: list[str] | None = None, + included_files: list[str] | None = None, + ): + """ + Prepare the retriever for a repository. + Will load database from local storage if available. + + Args: + repo_url_or_path: URL or local path to the repository + access_token: Optional access token for private repositories + excluded_dirs: Optional list of directories to exclude from processing + excluded_files: Optional list of file patterns to exclude from processing + included_dirs: Optional list of directories to include exclusively + included_files: Optional list of file patterns to include exclusively + """ + self.initialize_db_manager() + self.repo_url_or_path = repo_url_or_path + self.transformed_docs = self.db_manager.prepare_database( + repo_url_or_path, + type, + access_token, + embedder_type=self.embedder_type, + excluded_dirs=excluded_dirs, + excluded_files=excluded_files, + included_dirs=included_dirs, + included_files=included_files, + ) + logger.info(f"Loaded {len(self.transformed_docs)} documents for retrieval") + + # Validate and filter embeddings to ensure consistent sizes + self.transformed_docs = self._validate_and_filter_embeddings( + self.transformed_docs + ) + + if not self.transformed_docs: + raise ValueError( + "No valid documents with embeddings found. Cannot create retriever." + ) + + logger.info( + f"Using {len(self.transformed_docs)} documents with valid embeddings for retrieval" + ) + + try: + # Use the appropriate embedder for retrieval + self.retriever = FAISSRetriever( + **configs["retriever"], + embedder=self.embedder, + documents=self.transformed_docs, + document_map_func=lambda doc: doc.vector, + ) + logger.info("FAISS retriever created successfully") + except Exception as e: + logger.error(f"Error creating FAISS retriever: {str(e)}") + # Try to provide more specific error information + if "All embeddings should be of the same size" in str(e): + logger.error( + "Embedding size validation failed. This suggests there are still inconsistent embedding sizes." + ) + # Log embedding sizes for debugging + sizes = [] + for i, doc in enumerate( + self.transformed_docs[:10] + ): # Check first 10 docs + if hasattr(doc, "vector") and doc.vector is not None: + try: + size = _get_document_vector_size(doc) or "unknown" + sizes.append(f"doc_{i}: {size}") + except Exception: + sizes.append(f"doc_{i}: error") + logger.error(f"Sample embedding sizes: {', '.join(sizes)}") + raise + + async def aprepare_retriever( + self, + repo_url_or_path: str, + type: str = "github", + access_token: str | None = None, + excluded_dirs: list[str] | None = None, + excluded_files: list[str] | None = None, + included_dirs: list[str] | None = None, + included_files: list[str] | None = None, + ): + """Async version of the original `prepare_retriever`. + + Reuse the synchronous `prepare_retriever` implementation, but runs it in + a worker thread via `asyncio.to_thread` so that blocking operations (such + as git.clone, file io, embedding calls) do not stall the outer event loop. + Concurrency is bounded by a module-level semaphore, set by system variable + 'DEEPWIKI_MAX_CONCURRENT_RAG'. + + Args: + repo_url_or_path: URL or local path to the repository + access_token: Optional access token for private repositories + excluded_dirs: Optional list of directories to exclude from processing + excluded_files: Optional list of file patterns to exclude from processing + included_dirs: Optional list of directories to include exclusively + included_files: Optional list of file patterns to include exclusively + """ + async with _get_rag_semaphore(): + return await asyncio.to_thread( + self.prepare_retriever, + repo_url_or_path, + type=type, + access_token=access_token, + excluded_dirs=excluded_dirs, + excluded_files=excluded_files, + included_dirs=included_dirs, + included_files=included_files, + ) + + def call(self, query: str, language: str = "en") -> tuple[list]: + """ + Process a query using RAG. + + Args: + query: The user's query + + Returns: + Tuple of (RAGAnswer, retrieved_documents) + """ + try: + retrieved_documents = self.retriever(query) + + # Fill in the documents + retrieved_documents[0].documents = [ + self.transformed_docs[doc_index] + for doc_index in retrieved_documents[0].doc_indices + ] + + return retrieved_documents + + except Exception: + logger.exception("Error in RAG call.") + + # Create error response + error_response = RAGAnswer( + rationale="Error occurred while processing the query.", + answer="I apologize, but I encountered an error while processing your question. Please try again or rephrase your question.", + ) + return error_response, [] + + async def acall(self, query: str, language: str = "en") -> tuple[list]: + """Async version of the original `call` method.""" + return await asyncio.to_thread(self.call, query, language) diff --git a/api/repository.py b/api/repository.py new file mode 100644 index 000000000..aab174c41 --- /dev/null +++ b/api/repository.py @@ -0,0 +1,306 @@ +import base64 +import json +from urllib.parse import quote, urlparse + +import requests +from requests.exceptions import RequestException + +from api.logger import get_logger + +logger = get_logger(__name__) + + +def _get_github_file_content( + repo_url: str, file_path: str, access_token: str = None +) -> str: + """ + Retrieves the content of a file from a GitHub repository using the GitHub API. + Supports both public GitHub (github.com) and GitHub Enterprise (custom domains). + + Args: + repo_url (str): The URL of the GitHub repository + (e.g., "https://github.com/username/repo" or "https://github.company.com/username/repo") + file_path (str): The path to the file within the repository (e.g., "src/main.py") + access_token (str, optional): GitHub personal access token for private repositories + + Returns: + str: The content of the file as a string + + Raises: + ValueError: If the file cannot be fetched or if the URL is not a valid GitHub URL + """ + try: + # Parse the repository URL to support both github.com and enterprise GitHub + parsed_url = urlparse(repo_url) + if not parsed_url.scheme or not parsed_url.netloc: + raise ValueError("Not a valid GitHub repository URL") + + # Check if it's a GitHub-like URL structure + path_parts = parsed_url.path.strip("/").split("/") + if len(path_parts) < 2: + raise ValueError( + "Invalid GitHub URL format - expected format: https://domain/owner/repo" + ) + + owner = path_parts[-2] + repo = path_parts[-1].replace(".git", "") + + # Determine the API base URL + if parsed_url.netloc == "github.com": + # Public GitHub + api_base = "https://api.github.com" + else: + # GitHub Enterprise - API is typically at https://domain/api/v3/ + api_base = f"{parsed_url.scheme}://{parsed_url.netloc}/api/v3" + + # Use GitHub API to get file content + # The API endpoint for getting file content is: /repos/{owner}/{repo}/contents/{path} + api_url = f"{api_base}/repos/{owner}/{repo}/contents/{file_path}" + + # Fetch file content from GitHub API + headers = {} + if access_token: + headers["Authorization"] = f"token {access_token}" + logger.info(f"Fetching file content from GitHub API: {api_url}") + try: + response = requests.get(api_url, headers=headers) + response.raise_for_status() + except RequestException as e: + raise ValueError(f"Error fetching file content: {e}") + try: + content_data = response.json() + except json.JSONDecodeError: + raise ValueError("Invalid response from GitHub API") + + # Check if we got an error response + if "message" in content_data and "documentation_url" in content_data: + raise ValueError(f"GitHub API error: {content_data['message']}") + + # GitHub API returns file content as base64 encoded string + if "content" in content_data and "encoding" in content_data: + if content_data["encoding"] == "base64": + # The content might be split into lines, so join them first + content_base64 = content_data["content"].replace("\n", "") + content = base64.b64decode(content_base64).decode("utf-8") + return content + else: + raise ValueError(f"Unexpected encoding: {content_data['encoding']}") + else: + raise ValueError("File content not found in GitHub API response") + + except Exception as e: + raise ValueError(f"Failed to get file content: {str(e)}") + + +def _get_gitlab_file_content( + repo_url: str, file_path: str, access_token: str = None +) -> str: + """ + Retrieves the content of a file from a GitLab repository (cloud or self-hosted). + + Args: + repo_url (str): The GitLab repo URL (e.g., "https://gitlab.com/username/repo" or "http://localhost/group/project") + file_path (str): File path within the repository (e.g., "src/main.py") + access_token (str, optional): GitLab personal access token + + Returns: + str: File content + + Raises: + ValueError: If anything fails + """ + try: + # Parse and validate the URL + parsed_url = urlparse(repo_url) + if not parsed_url.scheme or not parsed_url.netloc: + raise ValueError("Not a valid GitLab repository URL") + + gitlab_domain = f"{parsed_url.scheme}://{parsed_url.netloc}" + if parsed_url.port not in (None, 80, 443): + gitlab_domain += f":{parsed_url.port}" + path_parts = parsed_url.path.strip("/").split("/") + if len(path_parts) < 2: + raise ValueError( + "Invalid GitLab URL format — expected something like https://gitlab.domain.com/group/project" + ) + + # Build project path and encode for API + project_path = "/".join(path_parts).replace(".git", "") + encoded_project_path = quote(project_path, safe="") + + # Encode file path + encoded_file_path = quote(file_path, safe="") + + # Try to get the default branch from the project info + default_branch = None + try: + project_info_url = f"{gitlab_domain}/api/v4/projects/{encoded_project_path}" + project_headers = {} + if access_token: + project_headers["PRIVATE-TOKEN"] = access_token + + project_response = requests.get(project_info_url, headers=project_headers) + if project_response.status_code == 200: + project_data = project_response.json() + default_branch = project_data.get("default_branch", "main") + logger.info(f"Found default branch: {default_branch}") + else: + logger.warning( + f"Could not fetch project info, using 'main' as default branch" + ) + default_branch = "main" + except Exception as e: + logger.warning( + f"Error fetching project info: {e}, using 'main' as default branch" + ) + default_branch = "main" + + api_url = f"{gitlab_domain}/api/v4/projects/{encoded_project_path}/repository/files/{encoded_file_path}/raw?ref={default_branch}" + # Fetch file content from GitLab API + headers = {} + if access_token: + headers["PRIVATE-TOKEN"] = access_token + logger.info(f"Fetching file content from GitLab API: {api_url}") + try: + response = requests.get(api_url, headers=headers) + response.raise_for_status() + content = response.text + except RequestException as e: + raise ValueError(f"Error fetching file content: {e}") + + # Check for GitLab error response (JSON instead of raw file) + if content.startswith("{") and '"message":' in content: + try: + error_data = json.loads(content) + if "message" in error_data: + raise ValueError(f"GitLab API error: {error_data['message']}") + except json.JSONDecodeError: + pass + + return content + + except Exception as e: + raise ValueError(f"Failed to get file content: {str(e)}") + + +def _get_bitbucket_file_content( + repo_url: str, file_path: str, access_token: str = None +) -> str: + """ + Retrieves the content of a file from a Bitbucket repository using the Bitbucket API. + + Args: + repo_url (str): The URL of the Bitbucket repository (e.g., "https://bitbucket.org/username/repo") + file_path (str): The path to the file within the repository (e.g., "src/main.py") + access_token (str, optional): Bitbucket personal access token for private repositories + + Returns: + str: The content of the file as a string + """ + try: + # Extract owner and repo name from Bitbucket URL + if not ( + repo_url.startswith("https://bitbucket.org/") + or repo_url.startswith("http://bitbucket.org/") + ): + raise ValueError("Not a valid Bitbucket repository URL") + + parts = repo_url.rstrip("/").split("/") + if len(parts) < 5: + raise ValueError("Invalid Bitbucket URL format") + + owner = parts[-2] + repo = parts[-1].replace(".git", "") + + # Try to get the default branch from the repository info + default_branch = None + try: + repo_info_url = f"https://api.bitbucket.org/2.0/repositories/{owner}/{repo}" + repo_headers = {} + if access_token: + repo_headers["Authorization"] = f"Bearer {access_token}" + + repo_response = requests.get(repo_info_url, headers=repo_headers) + if repo_response.status_code == 200: + repo_data = repo_response.json() + default_branch = repo_data.get("mainbranch", {}).get("name", "main") + logger.info(f"Found default branch: {default_branch}") + else: + logger.warning( + f"Could not fetch repository info, using 'main' as default branch" + ) + default_branch = "main" + except Exception as e: + logger.warning( + f"Error fetching repository info: {e}, using 'main' as default branch" + ) + default_branch = "main" + + # Use Bitbucket API to get file content + # The API endpoint for getting file content is: /2.0/repositories/{owner}/{repo}/src/{branch}/{path} + api_url = f"https://api.bitbucket.org/2.0/repositories/{owner}/{repo}/src/{default_branch}/{file_path}" + + # Fetch file content from Bitbucket API + headers = {} + if access_token: + headers["Authorization"] = f"Bearer {access_token}" + logger.info(f"Fetching file content from Bitbucket API: {api_url}") + try: + response = requests.get(api_url, headers=headers) + if response.status_code == 200: + content = response.text + elif response.status_code == 404: + raise ValueError( + "File not found on Bitbucket. Please check the file path and repository." + ) + elif response.status_code == 401: + raise ValueError( + "Unauthorized access to Bitbucket. Please check your access token." + ) + elif response.status_code == 403: + raise ValueError( + "Forbidden access to Bitbucket. You might not have permission to access this file." + ) + elif response.status_code == 500: + raise ValueError( + "Internal server error on Bitbucket. Please try again later." + ) + else: + response.raise_for_status() + content = response.text + return content + except RequestException as e: + raise ValueError(f"Error fetching file content: {e}") + + except Exception as e: + raise ValueError(f"Failed to get file content: {str(e)}") + + +def get_repo_content( + repo_url: str, file_path: str, repo_type: str = None, access_token: str = None +) -> str: + """ + Retrieves the content of a file from a Git repository (GitHub or GitLab). + + Args: + repo_type (str): Type of repository + repo_url (str): The URL of the repository + file_path (str): The path to the file within the repository + access_token (str, optional): Access token for private repositories + + Returns: + str: The content of the file as a string + + Raises: + ValueError: If the file cannot be fetched or if the URL is not valid + """ + if repo_type == "github": + return _get_github_file_content(repo_url, file_path, access_token) + elif repo_type == "gitlab": + return _get_gitlab_file_content(repo_url, file_path, access_token) + elif repo_type == "bitbucket": + return _get_bitbucket_file_content(repo_url, file_path, access_token) + else: + raise ValueError( + "Unsupported repository type. Only GitHub, GitLab, and Bitbucket are supported." + ) diff --git a/api/simple_chat.py b/api/simple_chat.py index a6e7e8aae..583765a6f 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -1,30 +1,32 @@ import asyncio import logging -from typing import Callable from functools import partial +from typing import Callable from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse -from api.chat import ChatStreamer, prompt_builder, is_token_limit_error -from api.config import get_model_config, configs -from api.data_pipeline import count_tokens, get_file_content -from api.rag import RAG, MAX_INPUT_TOKENS +from api.chat import ChatStreamer, is_token_limit_error, prompt_builder +from api.chat_model import ChatCompletionRequest +from api.config import configs, get_model_config +from api.logging_config import setup_logging from api.prompts import ( - DEEP_RESEARCH_FIRST_ITERATION_PROMPT, DEEP_RESEARCH_FINAL_ITERATION_PROMPT, + DEEP_RESEARCH_FIRST_ITERATION_PROMPT, DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT, - SIMPLE_CHAT_SYSTEM_PROMPT + SIMPLE_CHAT_SYSTEM_PROMPT, ) -from api.chat_model import ChatCompletionRequest +from api.rag import RAG, count_tokens +from api.repository import get_repo_content # Configure logging -from api.logging_config import setup_logging - setup_logging() logger = logging.getLogger(__name__) +# Maximum token limit for embedding models +MAX_INPUT_TOKENS = 7500 # Safe threshold below 8192 token limit + # Initialize FastAPI app app = FastAPI( @@ -254,7 +256,7 @@ async def chat_completions_stream(request: ChatCompletionRequest): if request.filePath: try: file_content = await asyncio.to_thread( - get_file_content, + get_repo_content, repo_url=request.repo_url, file_path=request.filePath, repo_type=request.type, diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index d72efc2b5..25458aff2 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -5,28 +5,32 @@ from fastapi import WebSocket, WebSocketDisconnect -from api.chat import ChatStreamer, prompt_builder, is_token_limit_error +from api.chat import ChatStreamer, is_token_limit_error, prompt_builder +from api.chat_model import ChatCompletionRequest from api.config import ( - get_model_config, configs, + get_model_config, ) -from api.data_pipeline import count_tokens, get_file_content -from api.rag import RAG, MAX_INPUT_TOKENS +from api.logging_config import setup_logging from api.prompts import ( - DEEP_RESEARCH_FIRST_ITERATION_PROMPT, DEEP_RESEARCH_FINAL_ITERATION_PROMPT, + DEEP_RESEARCH_FIRST_ITERATION_PROMPT, DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT, SIMPLE_CHAT_SYSTEM_PROMPT, ) -from api.chat_model import ChatCompletionRequest +from api.rag import RAG, count_tokens +from api.repository import get_repo_content # Configure logging -from api.logging_config import setup_logging - setup_logging() logger = logging.getLogger(__name__) +# Maximum token limit for embedding models +MAX_INPUT_TOKENS = 7500 # Safe threshold below 8192 token limit + + + async def handle_websocket_chat(websocket: WebSocket): """ Handle WebSocket connection for chat completions. @@ -251,7 +255,7 @@ async def handle_websocket_chat(websocket: WebSocket): if request.filePath: try: file_content = await asyncio.to_thread( - get_file_content, + get_repo_content, repo_url=request.repo_url, file_path=request.filePath, repo_type=request.type, From 94f43f0089cdcff31987c5cf56537f10f7af60ca Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 26 Jul 2026 22:37:17 +0800 Subject: [PATCH 2/7] Refactor dialog turns retrieval in call method Simplified the logic for retrieving dialog turns from the current conversation. --- api/rag/rag.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/api/rag/rag.py b/api/rag/rag.py index 069a84c49..2382b7622 100644 --- a/api/rag/rag.py +++ b/api/rag/rag.py @@ -86,13 +86,9 @@ def __init__(self): def call(self) -> dict: """Return the conversation history as a dictionary.""" - all_dialog_turns = ( - {} - if not self.current_conversation - else { - dialog_turn.id: dialog_turn for dialog_turn in self.current_conversation - } - ) + all_dialog_turns = { + dialog_turn.id: dialog_turn for dialog_turn in self.current_conversation + } logger.info(f"Returning {len(all_dialog_turns)} dialog turns from memory") return all_dialog_turns From 9b6d258572a378a2b3b5067fd59199e6e8c91403 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 27 Jul 2026 16:42:59 +0800 Subject: [PATCH 3/7] fix: fix rag return type, removing `RAGAnswer` dataclass object. --- api/rag/rag.py | 40 ++++++++++++---------------------------- 1 file changed, 12 insertions(+), 28 deletions(-) diff --git a/api/rag/rag.py b/api/rag/rag.py index 2382b7622..26f2eea93 100644 --- a/api/rag/rag.py +++ b/api/rag/rag.py @@ -3,12 +3,17 @@ import os from collections import defaultdict from collections.abc import Sized -from dataclasses import dataclass, field from uuid import uuid4 import adalflow as adal from adalflow.components.retriever.faiss_retriever import FAISSRetriever -from adalflow.core.types import AssistantResponse, DialogTurn, Document, UserQuery +from adalflow.core.types import ( + AssistantResponse, + DialogTurn, + Document, + RetrieverOutput, + UserQuery, +) from api.config import configs from api.rag.pipeline import DatabaseManager @@ -116,21 +121,6 @@ def add_dialog_turn(self, user_query: str, assistant_response: str) -> None: ) -@dataclass -class RAGAnswer(adal.DataClass): - rationale: str = field( - default="", metadata={"desc": "Chain of thoughts for the answer."} - ) - answer: str = field( - default="", - metadata={ - "desc": "Answer to the user query, formatted in markdown for beautiful rendering with react-markdown. DO NOT include ``` triple backticks fences at the beginning or end of your answer." - }, - ) - - __output_fields__ = ["rationale", "answer"] - - def _get_document_vector_size(document: Document) -> int | None: if hasattr(document.vector, "shape"): embedding_size = ( @@ -366,7 +356,7 @@ async def aprepare_retriever( included_files=included_files, ) - def call(self, query: str, language: str = "en") -> tuple[list]: + def call(self, query: str | list[str], language: str = "en") -> list[RetrieverOutput]: """ Process a query using RAG. @@ -374,7 +364,7 @@ def call(self, query: str, language: str = "en") -> tuple[list]: query: The user's query Returns: - Tuple of (RAGAnswer, retrieved_documents) + list of RetrieverOutput. """ try: retrieved_documents = self.retriever(query) @@ -388,15 +378,9 @@ def call(self, query: str, language: str = "en") -> tuple[list]: return retrieved_documents except Exception: - logger.exception("Error in RAG call.") - - # Create error response - error_response = RAGAnswer( - rationale="Error occurred while processing the query.", - answer="I apologize, but I encountered an error while processing your question. Please try again or rephrase your question.", - ) - return error_response, [] + logger.exception("Error in RAG call, returning empty list") + return [] - async def acall(self, query: str, language: str = "en") -> tuple[list]: + async def acall(self, query: str, language: str = "en") -> list[RetrieverOutput]: """Async version of the original `call` method.""" return await asyncio.to_thread(self.call, query, language) From b8a3cbebdd04667a787394ddd360912319db31b5 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Tue, 28 Jul 2026 09:30:27 +0800 Subject: [PATCH 4/7] add rag `_validate_and_filter` unittests --- tests/backend/rag/test_rag.py | 37 +++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 tests/backend/rag/test_rag.py diff --git a/tests/backend/rag/test_rag.py b/tests/backend/rag/test_rag.py new file mode 100644 index 000000000..f4a9399de --- /dev/null +++ b/tests/backend/rag/test_rag.py @@ -0,0 +1,37 @@ +import numpy as np +from adalflow.core.types import Document + +from api.rag import RAG + + +def test_rag_valid_filter_documents(): + doc_list = [ + Document( + text="test1", + vector=[10, 11, 12], + meta_data={}, + ), + Document( + text="test2", + vector=np.array([10, 11, 12]), + meta_data={}, + ), + Document( + text="test3", + vector=(10, 11, 12), + meta_data={}, + ), + Document( + text="invalid1", + vector=np.array([10, 11, 12, 13]), + meta_data={}, + ), + Document( + text="invalid2", + vector=None, + meta_data={}, + ), + ] + + validated_docs = doc_list.copy()[:3] + assert validated_docs == RAG._validate_and_filter_embeddings(doc_list) From 6e4c870caf04e926e296b7dbf652c8c2f1bb6453 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Tue, 28 Jul 2026 16:46:16 +0800 Subject: [PATCH 5/7] fix: fix `_should_process_file` implementation when dealing with file/directory prefix --- api/rag/pipeline.py | 425 ++++++++++++++++++++++++++------------------ 1 file changed, 252 insertions(+), 173 deletions(-) diff --git a/api/rag/pipeline.py b/api/rag/pipeline.py index 4a1c1b1dd..e65758816 100644 --- a/api/rag/pipeline.py +++ b/api/rag/pipeline.py @@ -11,7 +11,7 @@ from adalflow.core.types import Document, List from adalflow.utils import get_adalflow_default_root_path -from api.config import DEFAULT_EXCLUDED_DIRS, DEFAULT_EXCLUDED_FILES, configs +from api.config import configs from api.tools.embedder import get_embedder # Configure logging @@ -20,7 +20,10 @@ # Maximum token limit for OpenAI embedding models MAX_EMBEDDING_TOKENS = 8192 -def count_tokens(text: str, embedder_type: str = None, is_ollama_embedder: bool = None) -> int: + +def count_tokens( + text: str, embedder_type: str = None, is_ollama_embedder: bool = None +) -> int: """ Count the number of tokens in a text string using tiktoken. @@ -37,21 +40,22 @@ def count_tokens(text: str, embedder_type: str = None, is_ollama_embedder: bool try: # Handle backward compatibility if embedder_type is None and is_ollama_embedder is not None: - embedder_type = 'ollama' if is_ollama_embedder else None - + embedder_type = "ollama" if is_ollama_embedder else None + # Determine embedder type if not specified if embedder_type is None: from api.config import get_embedder_type + embedder_type = get_embedder_type() # Choose encoding based on embedder type - if embedder_type == 'ollama': + if embedder_type == "ollama": # Ollama typically uses cl100k_base encoding encoding = tiktoken.get_encoding("cl100k_base") - elif embedder_type == 'google': + elif embedder_type == "google": # Google uses similar tokenization to GPT models for rough estimation encoding = tiktoken.get_encoding("cl100k_base") - elif embedder_type == 'bedrock': + elif embedder_type == "bedrock": # Bedrock embedding models vary; use a common GPT-like encoding for rough estimation encoding = tiktoken.get_encoding("cl100k_base") else: # OpenAI or default @@ -65,7 +69,10 @@ def count_tokens(text: str, embedder_type: str = None, is_ollama_embedder: bool # Rough approximation: 4 characters per token return len(text) // 4 -def download_repo(repo_url: str, local_path: str, repo_type: str = None, access_token: str = None) -> str: + +def download_repo( + repo_url: str, local_path: str, repo_type: str = None, access_token: str = None +) -> str: """ Downloads a Git repository (GitHub, GitLab, or Bitbucket) to a specified local path. @@ -91,7 +98,9 @@ def download_repo(repo_url: str, local_path: str, repo_type: str = None, access_ # Check if repository already exists if os.path.exists(local_path) and os.listdir(local_path): # Directory exists and is not empty - logger.warning(f"Repository already exists at {local_path}. Using existing repository.") + logger.warning( + f"Repository already exists at {local_path}. Using existing repository." + ) return f"Using existing repository at {local_path}" # Ensure the local path exists @@ -102,15 +111,33 @@ def download_repo(repo_url: str, local_path: str, repo_type: str = None, access_ if access_token: parsed = urlparse(repo_url) # URL-encode the token to handle special characters - encoded_token = quote(access_token, safe='') + encoded_token = quote(access_token, safe="") # Determine the repository type and format the URL accordingly if repo_type == "github": # Format: https://{token}@{domain}/owner/repo.git # Works for both github.com and enterprise GitHub domains - clone_url = urlunparse((parsed.scheme, f"{encoded_token}@{parsed.netloc}", parsed.path, '', '', '')) + clone_url = urlunparse( + ( + parsed.scheme, + f"{encoded_token}@{parsed.netloc}", + parsed.path, + "", + "", + "", + ) + ) elif repo_type == "gitlab": # Format: https://oauth2:{token}@gitlab.com/owner/repo.git - clone_url = urlunparse((parsed.scheme, f"oauth2:{encoded_token}@{parsed.netloc}", parsed.path, '', '', '')) + clone_url = urlunparse( + ( + parsed.scheme, + f"oauth2:{encoded_token}@{parsed.netloc}", + parsed.path, + "", + "", + "", + ) + ) elif repo_type == "bitbucket": # Bitbucket has two token formats with different auth schemes: # - HTTP access tokens (prefix "ATCTT") use x-bitbucket-api-token-auth @@ -121,7 +148,16 @@ def download_repo(repo_url: str, local_path: str, repo_type: str = None, access_ else: auth_scheme = "x-token-auth" # Format: https://{auth_scheme}:{token}@bitbucket.org/owner/repo.git - clone_url = urlunparse((parsed.scheme, f"{auth_scheme}:{encoded_token}@{parsed.netloc}", parsed.path, '', '', '')) + clone_url = urlunparse( + ( + parsed.scheme, + f"{auth_scheme}:{encoded_token}@{parsed.netloc}", + parsed.path, + "", + "", + "", + ) + ) logger.info("Using access token for authentication") @@ -139,24 +175,107 @@ def download_repo(repo_url: str, local_path: str, repo_type: str = None, access_ return result.stdout.decode("utf-8") except subprocess.CalledProcessError as e: - error_msg = e.stderr.decode('utf-8') + error_msg = e.stderr.decode("utf-8") # Sanitize error message to remove any tokens (both raw and URL-encoded) if access_token: # Remove raw token error_msg = error_msg.replace(access_token, "***TOKEN***") # Also remove URL-encoded token to prevent leaking encoded version - encoded_token = quote(access_token, safe='') + encoded_token = quote(access_token, safe="") error_msg = error_msg.replace(encoded_token, "***TOKEN***") raise ValueError(f"Error during cloning: {error_msg}") except Exception as e: raise ValueError(f"An unexpected error occurred: {str(e)}") -# Alias for backward compatibility -download_github_repo = download_repo -def read_all_documents(path: str, embedder_type: str = None, is_ollama_embedder: bool = None, - excluded_dirs: List[str] = None, excluded_files: List[str] = None, - included_dirs: List[str] = None, included_files: List[str] = None): +def _should_process_file( + file_path: Path, + use_inclusion: bool, + included_dirs: list[str] | None, + included_files: list[str] | None, + excluded_dirs: list[str] | None, + excluded_files: list[str] | None, +) -> bool: + """ + Determine if a file should be processed based on inclusion/exclusion rules. + + Args: + file_path (str): The file path to check + use_inclusion (bool): Whether to use inclusion mode + included_dirs (List[str]): List of directories to include + included_files (List[str]): List of files to include + excluded_dirs (List[str]): List of directories to exclude + excluded_files (List[str]): List of files to exclude + + Returns: + bool: True if the file should be processed, False otherwise + """ + if isinstance(file_path, str): + file_path = Path(file_path) + file_path_parts = file_path.resolve().parts + file_name = file_path_parts[-1] + + if use_inclusion: + # Inclusion mode: file must be in included directories or match included files + is_included = False + + # Check if file is in an included directory + if included_dirs: + for included in included_dirs: + clean_included = included.removeprefix("./").rstrip("/") + if clean_included in file_path_parts: + is_included = True + break + + # Check if file matches included file patterns + if not is_included and included_files: + for included_file in included_files: + if file_name == included_file or file_name.endswith(included_file): + is_included = True + break + + # If no inclusion rules are specified for a category, allow all files from that category + if not included_dirs and not included_files: + is_included = True + elif not included_dirs and included_files: + # Only file patterns specified, allow all directories + pass # is_included is already set based on file patterns + elif included_dirs and not included_files: + # Only directory patterns specified, allow all files in included directories + pass # is_included is already set based on directory patterns + + return is_included + else: + # Exclusion mode: file must not be in excluded directories or match excluded files + is_excluded = False + + # Check if file is in an excluded directory + if excluded_dirs: + for excluded in excluded_dirs: + clean_excluded = excluded.removeprefix("./").rstrip("/") + if clean_excluded in file_path_parts: + is_excluded = True + break + + # Check if file matches excluded file patterns + if not is_excluded and excluded_files: + for excluded_file in excluded_files: + if file_name == excluded_file: + is_excluded = True + break + + return not is_excluded + + +def read_all_documents( + path: str, + embedder_type: str = None, + is_ollama_embedder: bool = None, + excluded_dirs: list[str] | None = None, + excluded_files: list[str] | None = None, + included_dirs: list[str] | None = None, + included_files: list[str] | None = None, +): """ Recursively reads all documents in a directory and its subdirectories. @@ -180,12 +299,11 @@ def read_all_documents(path: str, embedder_type: str = None, is_ollama_embedder: """ # Handle backward compatibility if embedder_type is None and is_ollama_embedder is not None: - embedder_type = 'ollama' if is_ollama_embedder else None + embedder_type = "ollama" if is_ollama_embedder else None documents = [] # File extensions to look for, prioritizing code files - code_extensions = [".py", ".js", ".ts", ".java", ".cpp", ".c", ".h", ".hpp", ".go", ".rs", - ".jsx", ".tsx", ".html", ".css", ".php", ".swift", ".cs"] - doc_extensions = [".md", ".txt", ".rst", ".json", ".yaml", ".yml"] + code_extensions = configs.get("code_extensions", []) + doc_extensions = configs.get("doc_extensions", []) # Determine filtering mode: inclusion or exclusion use_inclusion_mode = bool(included_dirs or included_files) @@ -195,7 +313,7 @@ def read_all_documents(path: str, embedder_type: str = None, is_ollama_embedder: included_dirs = list(set(included_dirs)) if included_dirs else list() included_files = list(set(included_files)) if included_files else list() - logger.info(f"Using inclusion mode") + logger.info("Using inclusion mode") logger.info(f"Included directories: {included_dirs}") logger.info(f"Included files: {included_files}") @@ -204,16 +322,9 @@ def read_all_documents(path: str, embedder_type: str = None, is_ollama_embedder: excluded_files = [] else: # Exclusion mode: use default exclusions plus any additional ones - final_excluded_dirs = set(DEFAULT_EXCLUDED_DIRS) - final_excluded_files = set(DEFAULT_EXCLUDED_FILES) - - # Add any additional excluded directories from config - if "file_filters" in configs and "excluded_dirs" in configs["file_filters"]: - final_excluded_dirs.update(configs["file_filters"]["excluded_dirs"]) - - # Add any additional excluded files from config - if "file_filters" in configs and "excluded_files" in configs["file_filters"]: - final_excluded_files.update(configs["file_filters"]["excluded_files"]) + file_filters = configs.get("file_filters", {}) + final_excluded_dirs: set[str] = set(file_filters.get("excluded_dirs", [])) + final_excluded_files: set[str] = set(file_filters.get("excluded_files", [])) # Add any explicitly provided excluded directories and files if excluded_dirs is not None: @@ -228,83 +339,14 @@ def read_all_documents(path: str, embedder_type: str = None, is_ollama_embedder: included_dirs = [] included_files = [] - logger.info(f"Using exclusion mode") + logger.info("Using exclusion mode") logger.info(f"Excluded directories: {excluded_dirs}") logger.info(f"Excluded files: {excluded_files}") logger.info(f"Reading documents from {path}") - def should_process_file(file_path: str, use_inclusion: bool, included_dirs: List[str], included_files: List[str], - excluded_dirs: List[str], excluded_files: List[str]) -> bool: - """ - Determine if a file should be processed based on inclusion/exclusion rules. - - Args: - file_path (str): The file path to check - use_inclusion (bool): Whether to use inclusion mode - included_dirs (List[str]): List of directories to include - included_files (List[str]): List of files to include - excluded_dirs (List[str]): List of directories to exclude - excluded_files (List[str]): List of files to exclude - - Returns: - bool: True if the file should be processed, False otherwise - """ - file_path_parts = os.path.normpath(file_path).split(os.sep) - file_name = os.path.basename(file_path) - - if use_inclusion: - # Inclusion mode: file must be in included directories or match included files - is_included = False - - # Check if file is in an included directory - if included_dirs: - for included in included_dirs: - clean_included = included.strip("./").rstrip("/") - if clean_included in file_path_parts: - is_included = True - break - - # Check if file matches included file patterns - if not is_included and included_files: - for included_file in included_files: - if file_name == included_file or file_name.endswith(included_file): - is_included = True - break - - # If no inclusion rules are specified for a category, allow all files from that category - if not included_dirs and not included_files: - is_included = True - elif not included_dirs and included_files: - # Only file patterns specified, allow all directories - pass # is_included is already set based on file patterns - elif included_dirs and not included_files: - # Only directory patterns specified, allow all files in included directories - pass # is_included is already set based on directory patterns - - return is_included - else: - # Exclusion mode: file must not be in excluded directories or match excluded files - is_excluded = False - - # Check if file is in an excluded directory - for excluded in excluded_dirs: - clean_excluded = excluded.strip("./").rstrip("/") - if clean_excluded in file_path_parts: - is_excluded = True - break - - # Check if file matches excluded file patterns - if not is_excluded: - for excluded_file in excluded_files: - if file_name == excluded_file: - is_excluded = True - break - - return not is_excluded - for file_path in filter( - lambda p: should_process_file( + lambda p: _should_process_file( p, use_inclusion=use_inclusion_mode, included_dirs=included_dirs, @@ -312,48 +354,50 @@ def should_process_file(file_path: str, use_inclusion: bool, included_dirs: List excluded_dirs=excluded_dirs, excluded_files=excluded_files, ), - filter( - lambda p: p.suffix.lower() in code_extensions + doc_extensions, - Path(path).rglob(pattern="**/*"), - ), + filter( + lambda p: p.suffix.lower() in code_extensions + doc_extensions, + Path(path).rglob(pattern="**/*"), + ), ): - try: - with open(file_path, "r", encoding="utf-8") as f: - content = f.read() - relative_path = os.path.relpath(file_path, path) - - # Check token count - token_count = count_tokens(content, embedder_type) - if token_count > MAX_EMBEDDING_TOKENS * 10: - logger.warning(f"Skipping large file {relative_path}: Token count ({token_count}) exceeds limit") - continue - - file_ext = file_path.suffix.lower() - is_code = file_ext in code_extensions - # Determine if this is an implementation file - if is_code: - is_implementation = ( - not relative_path.startswith("test_") - and not relative_path.startswith("app_") - and "test" not in relative_path.lower() - ) - else: - is_implementation = False - - doc = Document( - text=content, - meta_data={ - "file_path": relative_path, - "type": file_ext, - "is_code": is_code, - "is_implementation": is_implementation, - "title": relative_path, - "token_count": token_count, - }, + try: + with open(file_path, "r", encoding="utf-8") as f: + content = f.read() + relative_path = os.path.relpath(file_path, path) + + # Check token count + token_count = count_tokens(content, embedder_type) + if token_count > MAX_EMBEDDING_TOKENS * 10: + logger.warning( + f"Skipping large file {relative_path}: Token count ({token_count}) exceeds limit" ) - documents.append(doc) - except Exception as e: - logger.error(f"Error reading {file_path}: {e}") + continue + + file_ext = file_path.suffix.lower() + is_code = file_ext in code_extensions + # Determine if this is an implementation file + if is_code: + is_implementation = ( + not relative_path.startswith("test_") + and not relative_path.startswith("app_") + and "test" not in relative_path.lower() + ) + else: + is_implementation = False + + doc = Document( + text=content, + meta_data={ + "file_path": relative_path, + "type": file_ext, + "is_code": is_code, + "is_implementation": is_implementation, + "title": relative_path, + "token_count": token_count, + }, + ) + documents.append(doc) + except Exception as e: + logger.error(f"Error reading {file_path}: {e}") logger.info(f"Found {len(documents)} documents") return documents @@ -375,8 +419,8 @@ def prepare_data_pipeline(embedder_type: str = None, is_ollama_embedder: bool = # Handle backward compatibility if embedder_type is None and is_ollama_embedder is not None: - embedder_type = 'ollama' if is_ollama_embedder else None - + embedder_type = "ollama" if is_ollama_embedder else None + # Determine embedder type if not specified if embedder_type is None: embedder_type = get_embedder_type() @@ -387,17 +431,19 @@ def prepare_data_pipeline(embedder_type: str = None, is_ollama_embedder: bool = embedder = get_embedder(embedder_type=embedder_type) batch_size = embedder_config.get("batch_size", 500) - embedder_transformer = ToEmbeddings( - embedder=embedder, batch_size=batch_size - ) + embedder_transformer = ToEmbeddings(embedder=embedder, batch_size=batch_size) data_transformer = adal.Sequential( splitter, embedder_transformer ) # sequential will chain together splitter and embedder return data_transformer + def transform_documents_and_save_to_db( - documents: List[Document], db_path: str, embedder_type: str = None, is_ollama_embedder: bool = None + documents: List[Document], + db_path: str, + embedder_type: str = None, + is_ollama_embedder: bool = None, ) -> LocalDB: """ Transforms a list of documents and saves them to a local database. @@ -422,6 +468,7 @@ def transform_documents_and_save_to_db( db.save_state(filepath=db_path) return db + class DatabaseManager: """ Manages the creation, loading, transformation, and persistence of LocalDB instances. @@ -432,10 +479,18 @@ def __init__(self): self.repo_url_or_path = None self.repo_paths = None - def prepare_database(self, repo_url_or_path: str, repo_type: str = None, access_token: str = None, - embedder_type: str = None, is_ollama_embedder: bool = None, - excluded_dirs: List[str] = None, excluded_files: List[str] = None, - included_dirs: List[str] = None, included_files: List[str] = None) -> List[Document]: + def prepare_database( + self, + repo_url_or_path: str, + repo_type: str = None, + access_token: str = None, + embedder_type: str = None, + is_ollama_embedder: bool = None, + excluded_dirs: List[str] = None, + excluded_files: List[str] = None, + included_dirs: List[str] = None, + included_files: List[str] = None, + ) -> List[Document]: """ Create a new database from the repository. @@ -457,12 +512,17 @@ def prepare_database(self, repo_url_or_path: str, repo_type: str = None, access_ """ # Handle backward compatibility if embedder_type is None and is_ollama_embedder is not None: - embedder_type = 'ollama' if is_ollama_embedder else None - + embedder_type = "ollama" if is_ollama_embedder else None + self.reset_database() self._create_repo(repo_url_or_path, repo_type, access_token) - return self.prepare_db_index(embedder_type=embedder_type, excluded_dirs=excluded_dirs, excluded_files=excluded_files, - included_dirs=included_dirs, included_files=included_files) + return self.prepare_db_index( + embedder_type=embedder_type, + excluded_dirs=excluded_dirs, + excluded_files=excluded_files, + included_dirs=included_dirs, + included_files=included_files, + ) def reset_database(self): """ @@ -474,7 +534,7 @@ def reset_database(self): def _extract_repo_name_from_url(self, repo_url_or_path: str, repo_type: str) -> str: # Extract owner and repo name to create unique identifier - url_parts = repo_url_or_path.rstrip('/').split('/') + url_parts = repo_url_or_path.rstrip("/").split("/") if repo_type in ["github", "gitlab", "bitbucket"] and len(url_parts) >= 5: # GitHub URL format: https://github.com/owner/repo @@ -487,7 +547,9 @@ def _extract_repo_name_from_url(self, repo_url_or_path: str, repo_type: str) -> repo_name = url_parts[-1].replace(".git", "") return repo_name - def _create_repo(self, repo_url_or_path: str, repo_type: str = None, access_token: str = None) -> None: + def _create_repo( + self, repo_url_or_path: str, repo_type: str = None, access_token: str = None + ) -> None: """ Download and prepare all paths. Paths: @@ -504,14 +566,18 @@ def _create_repo(self, repo_url_or_path: str, repo_type: str = None, access_toke try: # Strip whitespace to handle URLs with leading/trailing spaces repo_url_or_path = repo_url_or_path.strip() - + root_path = get_adalflow_default_root_path() os.makedirs(root_path, exist_ok=True) # url - if repo_url_or_path.startswith("https://") or repo_url_or_path.startswith("http://"): + if repo_url_or_path.startswith("https://") or repo_url_or_path.startswith( + "http://" + ): # Extract the repository name from the URL - repo_name = self._extract_repo_name_from_url(repo_url_or_path, repo_type) + repo_name = self._extract_repo_name_from_url( + repo_url_or_path, repo_type + ) logger.info(f"Extracted repo name: {repo_name}") save_repo_dir = os.path.join(root_path, "repos", repo_name) @@ -519,9 +585,13 @@ def _create_repo(self, repo_url_or_path: str, repo_type: str = None, access_toke # Check if the repository directory already exists and is not empty if not (os.path.exists(save_repo_dir) and os.listdir(save_repo_dir)): # Only download if the repository doesn't exist or is empty - download_repo(repo_url_or_path, save_repo_dir, repo_type, access_token) + download_repo( + repo_url_or_path, save_repo_dir, repo_type, access_token + ) else: - logger.info(f"Repository already exists at {save_repo_dir}. Using existing repository.") + logger.info( + f"Repository already exists at {save_repo_dir}. Using existing repository." + ) else: # local path repo_name = os.path.basename(repo_url_or_path) save_repo_dir = repo_url_or_path @@ -541,9 +611,15 @@ def _create_repo(self, repo_url_or_path: str, repo_type: str = None, access_toke logger.error(f"Failed to create repository structure: {e}") raise - def prepare_db_index(self, embedder_type: str = None, is_ollama_embedder: bool = None, - excluded_dirs: List[str] = None, excluded_files: List[str] = None, - included_dirs: List[str] = None, included_files: List[str] = None) -> List[Document]: + def prepare_db_index( + self, + embedder_type: str = None, + is_ollama_embedder: bool = None, + excluded_dirs: List[str] = None, + excluded_files: List[str] = None, + included_dirs: List[str] = None, + included_files: List[str] = None, + ) -> List[Document]: """ Prepare the indexed database for the repository. @@ -560,6 +636,7 @@ def prepare_db_index(self, embedder_type: str = None, is_ollama_embedder: bool = Returns: List[Document]: List of Document objects """ + def _embedding_vector_length(doc: Document) -> int: vector = getattr(doc, "vector", None) if vector is None: @@ -577,7 +654,7 @@ def _embedding_vector_length(doc: Document) -> int: # Handle backward compatibility if embedder_type is None and is_ollama_embedder is not None: - embedder_type = 'ollama' if is_ollama_embedder else None + embedder_type = "ollama" if is_ollama_embedder else None # check the database if self.repo_paths and os.path.exists(self.repo_paths["save_db_file"]): logger.info("Loading existing database...") @@ -615,7 +692,7 @@ def _embedding_vector_length(doc: Document) -> int: excluded_dirs=excluded_dirs, excluded_files=excluded_files, included_dirs=included_dirs, - included_files=included_files + included_files=included_files, ) self.db = transform_documents_and_save_to_db( documents, self.repo_paths["save_db_file"], embedder_type=embedder_type @@ -625,7 +702,9 @@ def _embedding_vector_length(doc: Document) -> int: logger.info(f"Total transformed documents: {len(transformed_docs)}") return transformed_docs - def prepare_retriever(self, repo_url_or_path: str, repo_type: str = None, access_token: str = None): + def prepare_retriever( + self, repo_url_or_path: str, repo_type: str = None, access_token: str = None + ): """ Prepare the retriever for a repository. This is a compatibility method for the isolated API. From 2e426c4020c103ecc03cc23f6cc291e56eefd058 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Tue, 28 Jul 2026 16:50:32 +0800 Subject: [PATCH 6/7] use file_filters in `repo.json` --- api/config.py | 43 +------------------------------------------ 1 file changed, 1 insertion(+), 42 deletions(-) diff --git a/api/config.py b/api/config.py index bd55b7247..fc7881ad7 100644 --- a/api/config.py +++ b/api/config.py @@ -293,47 +293,6 @@ def load_lang_config(): return loaded_config -# Default excluded directories and files -DEFAULT_EXCLUDED_DIRS: List[str] = [ - # Virtual environments and package managers - "./.venv/", "./venv/", "./env/", "./virtualenv/", - "./node_modules/", "./bower_components/", "./jspm_packages/", - # Version control - "./.git/", "./.svn/", "./.hg/", "./.bzr/", - # Cache and compiled files - "./__pycache__/", "./.pytest_cache/", "./.mypy_cache/", "./.ruff_cache/", "./.coverage/", - # Build and distribution - "./dist/", "./build/", "./out/", "./target/", "./bin/", "./obj/", - # Documentation - "./docs/", "./_docs/", "./site-docs/", "./_site/", - # IDE specific - "./.idea/", "./.vscode/", "./.vs/", "./.eclipse/", "./.settings/", - # Logs and temporary files - "./logs/", "./log/", "./tmp/", "./temp/", -] - -DEFAULT_EXCLUDED_FILES: List[str] = [ - "yarn.lock", "pnpm-lock.yaml", "npm-shrinkwrap.json", "poetry.lock", - "Pipfile.lock", "requirements.txt.lock", "Cargo.lock", "composer.lock", - ".lock", ".DS_Store", "Thumbs.db", "desktop.ini", "*.lnk", ".env", - ".env.*", "*.env", "*.cfg", "*.ini", ".flaskenv", ".gitignore", - ".gitattributes", ".gitmodules", ".github", ".gitlab-ci.yml", - ".prettierrc", ".eslintrc", ".eslintignore", ".stylelintrc", - ".editorconfig", ".jshintrc", ".pylintrc", ".flake8", "mypy.ini", - "pyproject.toml", "tsconfig.json", "webpack.config.js", "babel.config.js", - "rollup.config.js", "jest.config.js", "karma.conf.js", "vite.config.js", - "next.config.js", "*.min.js", "*.min.css", "*.bundle.js", "*.bundle.css", - "*.map", "*.gz", "*.zip", "*.tar", "*.tgz", "*.rar", "*.7z", "*.iso", - "*.dmg", "*.img", "*.msix", "*.appx", "*.appxbundle", "*.xap", "*.ipa", - "*.deb", "*.rpm", "*.msi", "*.exe", "*.dll", "*.so", "*.dylib", "*.o", - "*.obj", "*.jar", "*.war", "*.ear", "*.jsm", "*.class", "*.pyc", "*.pyd", - "*.pyo", "__pycache__", "*.a", "*.lib", "*.lo", "*.la", "*.slo", "*.dSYM", - "*.egg", "*.egg-info", "*.dist-info", "*.eggs", "node_modules", - "bower_components", "jspm_packages", "lib-cov", "coverage", "htmlcov", - ".nyc_output", ".tox", "dist", "build", "bld", "out", "bin", "target", - "packages/*/dist", "packages/*/build", ".output" -] - # Initialize empty configuration configs = {} @@ -356,7 +315,7 @@ def load_lang_config(): # Update repository configuration if repo_config: - for key in ["file_filters", "repository"]: + for key in ["file_filters", "repository", "code_extensions", "doc_extensions"]: if key in repo_config: configs[key] = repo_config[key] From 78e375a03e0466f5c6ec786e1f044de2de7046c6 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Tue, 28 Jul 2026 18:47:05 +0800 Subject: [PATCH 7/7] ruff formating --- api/config.py | 139 ++++++++++++++++++++++++++++++------------------- api/rag/rag.py | 8 ++- 2 files changed, 90 insertions(+), 57 deletions(-) diff --git a/api/config.py b/api/config.py index fc7881ad7..8172bc962 100644 --- a/api/config.py +++ b/api/config.py @@ -1,11 +1,9 @@ -import os import json import logging +import os import re from pathlib import Path -from typing import List, Union, Dict, Any - -logger = logging.getLogger(__name__) +from typing import Any, Dict, List, Union from api.clients import ( AzureAIClient, @@ -19,16 +17,18 @@ OpenRouterClient, ) +logger = logging.getLogger(__name__) + # Get API keys from environment variables -OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY') -LITELLM_API_KEY = os.environ.get('LITELLM_API_KEY') -GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY') -OPENROUTER_API_KEY = os.environ.get('OPENROUTER_API_KEY') -AWS_ACCESS_KEY_ID = os.environ.get('AWS_ACCESS_KEY_ID') -AWS_SECRET_ACCESS_KEY = os.environ.get('AWS_SECRET_ACCESS_KEY') -AWS_SESSION_TOKEN = os.environ.get('AWS_SESSION_TOKEN') -AWS_REGION = os.environ.get('AWS_REGION') -AWS_ROLE_ARN = os.environ.get('AWS_ROLE_ARN') +OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") +LITELLM_API_KEY = os.environ.get("LITELLM_API_KEY") +GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY") +OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY") +AWS_ACCESS_KEY_ID = os.environ.get("AWS_ACCESS_KEY_ID") +AWS_SECRET_ACCESS_KEY = os.environ.get("AWS_SECRET_ACCESS_KEY") +AWS_SESSION_TOKEN = os.environ.get("AWS_SESSION_TOKEN") +AWS_REGION = os.environ.get("AWS_REGION") +AWS_ROLE_ARN = os.environ.get("AWS_ROLE_ARN") # Set keys in environment (in case they're needed elsewhere in the code) if OPENAI_API_KEY: @@ -51,30 +51,45 @@ os.environ["AWS_ROLE_ARN"] = AWS_ROLE_ARN # Wiki authentication settings -raw_auth_mode = os.environ.get('DEEPWIKI_AUTH_MODE', 'False') -WIKI_AUTH_MODE = raw_auth_mode.lower() in ['true', '1', 't'] -WIKI_AUTH_CODE = os.environ.get('DEEPWIKI_AUTH_CODE', '') +raw_auth_mode = os.environ.get("DEEPWIKI_AUTH_MODE", "False") +WIKI_AUTH_MODE = raw_auth_mode.lower() in ["true", "1", "t"] +WIKI_AUTH_CODE = os.environ.get("DEEPWIKI_AUTH_CODE", "") # Embedder settings -EMBEDDER_TYPE = os.environ.get('DEEPWIKI_EMBEDDER_TYPE', 'openai').lower() +EMBEDDER_TYPE = os.environ.get("DEEPWIKI_EMBEDDER_TYPE", "openai").lower() # Get configuration directory from environment variable, or use default if not set -CONFIG_DIR = os.environ.get('DEEPWIKI_CONFIG_DIR', None) +CONFIG_DIR = os.environ.get("DEEPWIKI_CONFIG_DIR", None) # Client class mapping CLIENT_CLASSES = { - "GoogleGenAIClient": GoogleGenAIClient, - "GoogleEmbedderClient": GoogleEmbedderClient, - "OpenAIClient": OpenAIClient, - "LiteLLMClient" : LiteLLMClient, - "OpenRouterClient": OpenRouterClient, - "OllamaClient": OllamaClient, - "BedrockClient": BedrockClient, - "AzureAIClient": AzureAIClient, - "DashscopeClient": DashscopeClient + GoogleGenAIClient.__name__: GoogleGenAIClient, + GoogleEmbedderClient.__name__: GoogleEmbedderClient, + OpenAIClient.__name__: OpenAIClient, + LiteLLMClient.__name__: LiteLLMClient, + OpenRouterClient.__name__: OpenRouterClient, + OllamaClient.__name__: OllamaClient, + BedrockClient.__name__: BedrockClient, + AzureAIClient.__name__: AzureAIClient, + DashscopeClient.__name__: DashscopeClient, } -def replace_env_placeholders(config: Union[Dict[str, Any], List[Any], str, Any]) -> Union[Dict[str, Any], List[Any], str, Any]: + +_DEFAULT_PROVIDER_MAP = { + "google": GoogleGenAIClient, + "openai": OpenAIClient, + "litellm": LiteLLMClient, + "openrouter": OpenRouterClient, + "ollama": OllamaClient, + "bedrock": BedrockClient, + "azure": AzureAIClient, + "dashscope": DashscopeClient, +} + + +def replace_env_placeholders( + config: Union[Dict[str, Any], List[Any], str, Any], +) -> Union[Dict[str, Any], List[Any], str, Any]: """ Recursively replace placeholders like "${ENV_VAR}" in string values within a nested configuration structure (dicts, lists, strings) @@ -104,6 +119,7 @@ def replacer(match: re.Match[str]) -> str: # Handles numbers, booleans, None, etc. return config + # Load JSON configuration file def load_json_config(filename): try: @@ -120,7 +136,7 @@ def load_json_config(filename): logger.warning(f"Configuration file {config_path} does not exist") return {} - with open(config_path, 'r', encoding='utf-8') as f: + with open(config_path, "r", encoding="utf-8") as f: config = json.load(f) config = replace_env_placeholders(config) return config @@ -128,6 +144,7 @@ def load_json_config(filename): logger.error(f"Error loading configuration file {filename}: {str(e)}") return {} + # Load generator model configuration def load_generator_config(): generator_config = load_json_config("generator.json") @@ -137,25 +154,18 @@ def load_generator_config(): for provider_id, provider_config in generator_config["providers"].items(): # Try to set client class from client_class if provider_config.get("client_class") in CLIENT_CLASSES: - provider_config["model_client"] = CLIENT_CLASSES[provider_config["client_class"]] + provider_config["model_client"] = CLIENT_CLASSES[ + provider_config["client_class"] + ] # Fall back to default mapping based on provider_id - elif provider_id in ["google", "openai", "openrouter", "ollama", "bedrock", "azure", "dashscope", "litellm"]: - default_map = { - "google": GoogleGenAIClient, - "openai": OpenAIClient, - "litellm": LiteLLMClient, - "openrouter": OpenRouterClient, - "ollama": OllamaClient, - "bedrock": BedrockClient, - "azure": AzureAIClient, - "dashscope": DashscopeClient - } - provider_config["model_client"] = default_map[provider_id] + elif provider_id in _DEFAULT_PROVIDER_MAP: + provider_config["model_client"] = _DEFAULT_PROVIDER_MAP[provider_id] else: logger.warning(f"Unknown provider or client class: {provider_id}") return generator_config + # Load embedder configuration def load_embedder_config(): embedder_config = load_json_config("embedder.json") @@ -169,6 +179,7 @@ def load_embedder_config(): return embedder_config + def get_embedder_config(): """ Get the current embedder configuration based on DEEPWIKI_EMBEDDER_TYPE. @@ -177,15 +188,16 @@ def get_embedder_config(): dict: The embedder configuration with model_client resolved """ embedder_type = EMBEDDER_TYPE - if embedder_type == 'bedrock' and 'embedder_bedrock' in configs: + if embedder_type == "bedrock" and "embedder_bedrock" in configs: return configs.get("embedder_bedrock", {}) - elif embedder_type == 'google' and 'embedder_google' in configs: + elif embedder_type == "google" and "embedder_google" in configs: return configs.get("embedder_google", {}) - elif embedder_type == 'ollama' and 'embedder_ollama' in configs: + elif embedder_type == "ollama" and "embedder_ollama" in configs: return configs.get("embedder_ollama", {}) else: return configs.get("embedder", {}) + def is_ollama_embedder(): """ Check if the current embedder configuration uses OllamaClient. @@ -206,6 +218,7 @@ def is_ollama_embedder(): client_class = embedder_config.get("client_class", "") return client_class == "OllamaClient" + def is_google_embedder(): """ Check if the current embedder configuration uses GoogleEmbedderClient. @@ -226,6 +239,7 @@ def is_google_embedder(): client_class = embedder_config.get("client_class", "") return client_class == "GoogleEmbedderClient" + def is_bedrock_embedder(): """ Check if the current embedder configuration uses BedrockClient. @@ -244,26 +258,29 @@ def is_bedrock_embedder(): client_class = embedder_config.get("client_class", "") return client_class == "BedrockClient" + def get_embedder_type(): """ Get the current embedder type based on configuration. - + Returns: str: 'bedrock', 'ollama', 'google', or 'openai' (default) """ if is_bedrock_embedder(): - return 'bedrock' + return "bedrock" elif is_ollama_embedder(): - return 'ollama' + return "ollama" elif is_google_embedder(): - return 'google' + return "google" else: - return 'openai' + return "openai" + # Load repository and file filters configuration def load_repo_config(): return load_json_config("repo.json") + # Load language configuration def load_lang_config(): default_config = { @@ -277,22 +294,27 @@ def load_lang_config(): "vi": "Vietnamese (Tiếng Việt)", "pt-br": "Brazilian Portuguese (Português Brasileiro)", "fr": "Français (French)", - "ru": "Русский (Russian)" + "ru": "Русский (Russian)", }, - "default": "en" + "default": "en", } - loaded_config = load_json_config("lang.json") # Let load_json_config handle path and loading + loaded_config = load_json_config( + "lang.json" + ) # Let load_json_config handle path and loading if not loaded_config: return default_config if "supported_languages" not in loaded_config or "default" not in loaded_config: - logger.warning("Language configuration file 'lang.json' is malformed. Using default language configuration.") + logger.warning( + "Language configuration file 'lang.json' is malformed. Using default language configuration." + ) return default_config return loaded_config + # Initialize empty configuration configs = {} @@ -309,7 +331,14 @@ def load_lang_config(): # Update embedder configuration if embedder_config: - for key in ["embedder", "embedder_ollama", "embedder_google", "embedder_bedrock", "retriever", "text_splitter"]: + for key in [ + "embedder", + "embedder_ollama", + "embedder_google", + "embedder_bedrock", + "retriever", + "text_splitter", + ]: if key in embedder_config: configs[key] = embedder_config[key] diff --git a/api/rag/rag.py b/api/rag/rag.py index 26f2eea93..97fe4e91c 100644 --- a/api/rag/rag.py +++ b/api/rag/rag.py @@ -229,7 +229,9 @@ def _validate_and_filter_embeddings(documents: list[Document]) -> list: ) if not valid_documents: - logger.warning("No documents with valid embeddings remained after filtering") + logger.warning( + "No documents with valid embeddings remained after filtering" + ) else: logger.info( "Embedding validation complete: %d/%d documents have valid embeddings.", @@ -356,7 +358,9 @@ async def aprepare_retriever( included_files=included_files, ) - def call(self, query: str | list[str], language: str = "en") -> list[RetrieverOutput]: + def call( + self, query: str | list[str], language: str = "en" + ) -> list[RetrieverOutput]: """ Process a query using RAG.