diff --git a/api/config.py b/api/config.py index bd55b7247..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,62 +294,26 @@ 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 -# 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 = {} @@ -350,13 +331,20 @@ 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] # 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] diff --git a/api/data_pipeline.py b/api/data_pipeline.py deleted file mode 100644 index 16d15a84d..000000000 --- a/api/data_pipeline.py +++ /dev/null @@ -1,905 +0,0 @@ -import adalflow as adal -from adalflow.core.types import Document, List -from adalflow.components.data_process import TextSplitter, ToEmbeddings -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 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 api.tools.embedder import get_embedder - -# Configure logging -logger = logging.getLogger(__name__) - -# 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: - """ - Count the number of tokens in a text string using tiktoken. - - Args: - text (str): The text to count tokens for. - embedder_type (str, optional): The embedder type ('openai', 'google', 'ollama', 'bedrock'). - If None, will be determined from configuration. - is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. - If None, will be determined from configuration. - - Returns: - int: The number of tokens in the text. - """ - 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 - - # 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': - # Ollama typically uses cl100k_base encoding - encoding = tiktoken.get_encoding("cl100k_base") - elif embedder_type == 'google': - # Google uses similar tokenization to GPT models for rough estimation - encoding = tiktoken.get_encoding("cl100k_base") - 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 - # Use OpenAI embedding model encoding - encoding = tiktoken.encoding_for_model("text-embedding-3-small") - - return len(encoding.encode(text)) - except Exception as e: - # Fallback to a simple approximation if tiktoken fails - logger.warning(f"Error counting tokens with tiktoken: {e}") - # 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: - """ - Downloads a Git repository (GitHub, GitLab, or Bitbucket) to a specified local path. - - Args: - repo_type(str): Type of repository - repo_url (str): The URL of the Git repository to clone. - local_path (str): The local directory where the repository will be cloned. - access_token (str, optional): Access token for private repositories. - - Returns: - str: The output message from the `git` command. - """ - try: - # Check if Git is installed - logger.info(f"Preparing to clone repository to {local_path}") - subprocess.run( - ["git", "--version"], - check=True, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - - # 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.") - return f"Using existing repository at {local_path}" - - # Ensure the local path exists - os.makedirs(local_path, exist_ok=True) - - # Prepare the clone URL with access token if provided - clone_url = repo_url - if access_token: - parsed = urlparse(repo_url) - # URL-encode the token to handle special characters - 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, '', '', '')) - 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, '', '', '')) - elif repo_type == "bitbucket": - # Bitbucket has two token formats with different auth schemes: - # - HTTP access tokens (prefix "ATCTT") use x-bitbucket-api-token-auth - # - App passwords (deprecated, EOL June 2026) use x-token-auth - # Detect by token prefix so existing app password users keep working. - if access_token.startswith("ATCTT"): - auth_scheme = "x-bitbucket-api-token-auth" - 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, '', '', '')) - - logger.info("Using access token for authentication") - - # Clone the repository - logger.info(f"Cloning repository from {repo_url} to {local_path}") - # We use repo_url in the log to avoid exposing the token in logs - result = subprocess.run( - ["git", "clone", "--depth=1", "--single-branch", clone_url, local_path], - check=True, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - - logger.info("Repository cloned successfully") - return result.stdout.decode("utf-8") - - except subprocess.CalledProcessError as e: - 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='') - 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): - """ - Recursively reads all documents in a directory and its subdirectories. - - Args: - path (str): The root directory path. - embedder_type (str, optional): The embedder type ('openai', 'google', 'ollama'). - If None, will be determined from configuration. - is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. - If None, will be determined from configuration. - excluded_dirs (List[str], optional): List of directories to exclude from processing. - Overrides the default configuration if provided. - excluded_files (List[str], optional): List of file patterns to exclude from processing. - Overrides the default configuration if provided. - included_dirs (List[str], optional): List of directories to include exclusively. - When provided, only files in these directories will be processed. - included_files (List[str], optional): List of file patterns to include exclusively. - When provided, only files matching these patterns will be processed. - - Returns: - list: A list of Document objects with metadata. - """ - # Handle backward compatibility - if embedder_type is None and is_ollama_embedder is not 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"] - - # Determine filtering mode: inclusion or exclusion - use_inclusion_mode = bool(included_dirs or included_files) - - if use_inclusion_mode: - # Inclusion mode: only process specified directories and files - 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(f"Included directories: {included_dirs}") - logger.info(f"Included files: {included_files}") - - # Convert to lists for processing - excluded_dirs = [] - 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"]) - - # Add any explicitly provided excluded directories and files - if excluded_dirs is not None: - final_excluded_dirs.update(excluded_dirs) - - if excluded_files is not None: - final_excluded_files.update(excluded_files) - - # Convert back to lists for compatibility - excluded_dirs = list(final_excluded_dirs) - excluded_files = list(final_excluded_files) - included_dirs = [] - included_files = [] - - logger.info(f"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( - p, - use_inclusion=use_inclusion_mode, - included_dirs=included_dirs, - included_files=included_files, - excluded_dirs=excluded_dirs, - excluded_files=excluded_files, - ), - 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, - }, - ) - documents.append(doc) - except Exception as e: - logger.error(f"Error reading {file_path}: {e}") - - logger.info(f"Found {len(documents)} documents") - return documents - -def prepare_data_pipeline(embedder_type: str = None, is_ollama_embedder: bool = None): - """ - Creates and returns the data transformation pipeline. - - Args: - embedder_type (str, optional): The embedder type ('openai', 'google', 'ollama'). - If None, will be determined from configuration. - is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. - If None, will be determined from configuration. - - Returns: - adal.Sequential: The data transformation pipeline - """ - from api.config import get_embedder_config, get_embedder_type - - # Handle backward compatibility - if embedder_type is None and is_ollama_embedder is not 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() - - splitter = TextSplitter(**configs["text_splitter"]) - embedder_config = get_embedder_config() - - embedder = get_embedder(embedder_type=embedder_type) - - batch_size = embedder_config.get("batch_size", 500) - 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 -) -> LocalDB: - """ - Transforms a list of documents and saves them to a local database. - - Args: - documents (list): A list of `Document` objects. - db_path (str): The path to the local database file. - embedder_type (str, optional): The embedder type ('openai', 'google', 'ollama'). - If None, will be determined from configuration. - is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. - If None, will be determined from configuration. - """ - # Get the data transformer - data_transformer = prepare_data_pipeline(embedder_type, is_ollama_embedder) - - # Save the documents to a local database - db = LocalDB() - db.register_transformer(transformer=data_transformer, key="split_and_embed") - db.load(documents) - db.transform(key="split_and_embed") - os.makedirs(os.path.dirname(db_path), exist_ok=True) - 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. - """ - - def __init__(self): - self.db = None - 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]: - """ - Create a new database from the repository. - - Args: - repo_type(str): Type of repository - repo_url_or_path (str): The URL or local path of the repository - access_token (str, optional): Access token for private repositories - embedder_type (str, optional): Embedder type to use ('openai', 'google', 'ollama'). - If None, will be determined from configuration. - is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. - If None, will be determined from configuration. - excluded_dirs (List[str], optional): List of directories to exclude from processing - excluded_files (List[str], optional): List of file patterns to exclude from processing - included_dirs (List[str], optional): List of directories to include exclusively - included_files (List[str], optional): List of file patterns to include exclusively - - Returns: - List[Document]: List of Document objects - """ - # Handle backward compatibility - if embedder_type is None and is_ollama_embedder is not 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) - - def reset_database(self): - """ - Reset the database to its initial state. - """ - self.db = None - self.repo_url_or_path = None - self.repo_paths = None - - 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('/') - - if repo_type in ["github", "gitlab", "bitbucket"] and len(url_parts) >= 5: - # GitHub URL format: https://github.com/owner/repo - # GitLab URL format: https://gitlab.com/owner/repo or https://gitlab.com/group/subgroup/repo - # Bitbucket URL format: https://bitbucket.org/owner/repo - owner = url_parts[-2] - repo = url_parts[-1].replace(".git", "") - repo_name = f"{owner}_{repo}" - else: - 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: - """ - Download and prepare all paths. - Paths: - ~/.adalflow/repos/{owner}_{repo_name} (for url, local path will be the same) - ~/.adalflow/databases/{owner}_{repo_name}.pkl - - Args: - repo_type(str): Type of repository - repo_url_or_path (str): The URL or local path of the repository - access_token (str, optional): Access token for private repositories - """ - logger.info(f"Preparing repo storage for {repo_url_or_path}...") - - 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://"): - # Extract the repository name from the URL - 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) - - # 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) - else: - 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 - - save_db_file = os.path.join(root_path, "databases", f"{repo_name}.pkl") - os.makedirs(save_repo_dir, exist_ok=True) - os.makedirs(os.path.dirname(save_db_file), exist_ok=True) - - self.repo_paths = { - "save_repo_dir": save_repo_dir, - "save_db_file": save_db_file, - } - self.repo_url_or_path = repo_url_or_path - logger.info(f"Repo paths: {self.repo_paths}") - - except Exception as e: - 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]: - """ - Prepare the indexed database for the repository. - - Args: - embedder_type (str, optional): Embedder type to use ('openai', 'google', 'ollama'). - If None, will be determined from configuration. - is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. - If None, will be determined from configuration. - excluded_dirs (List[str], optional): List of directories to exclude from processing - excluded_files (List[str], optional): List of file patterns to exclude from processing - included_dirs (List[str], optional): List of directories to include exclusively - included_files (List[str], optional): List of file patterns to include exclusively - - Returns: - List[Document]: List of Document objects - """ - def _embedding_vector_length(doc: Document) -> int: - vector = getattr(doc, "vector", None) - if vector is None: - return 0 - try: - if hasattr(vector, "shape"): - if len(vector.shape) == 0: - return 0 - return int(vector.shape[-1]) - if hasattr(vector, "__len__"): - return int(len(vector)) - except Exception: - return 0 - return 0 - - # Handle backward compatibility - if embedder_type is None and is_ollama_embedder is not 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...") - try: - self.db = LocalDB.load_state(self.repo_paths["save_db_file"]) - documents = self.db.get_transformed_data(key="split_and_embed") - if documents: - lengths = [_embedding_vector_length(doc) for doc in documents] - non_empty = sum(1 for n in lengths if n > 0) - empty = len(lengths) - non_empty - sample_sizes = sorted({n for n in lengths if n > 0})[:3] - logger.info( - "Loaded %s documents from existing database (embeddings: %s non-empty, %s empty; sample_dims=%s)", - len(documents), - non_empty, - empty, - sample_sizes, - ) - - if non_empty == 0: - logger.warning( - "Existing database contains no usable embeddings. Rebuilding embeddings..." - ) - else: - return documents - except Exception as e: - logger.error(f"Error loading existing database: {e}") - # Continue to create a new database - - # prepare the database - logger.info("Creating new database...") - documents = read_all_documents( - self.repo_paths["save_repo_dir"], - embedder_type=embedder_type, - excluded_dirs=excluded_dirs, - excluded_files=excluded_files, - included_dirs=included_dirs, - included_files=included_files - ) - self.db = transform_documents_and_save_to_db( - documents, self.repo_paths["save_db_file"], embedder_type=embedder_type - ) - logger.info(f"Total documents: {len(documents)}") - transformed_docs = self.db.get_transformed_data(key="split_and_embed") - 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): - """ - Prepare the retriever for a repository. - This is a compatibility method for the isolated API. - - Args: - repo_type(str): Type of repository - repo_url_or_path (str): The URL or local path of the repository - access_token (str, optional): Access token for private repositories - - Returns: - List[Document]: List of Document objects - """ - return self.prepare_database(repo_url_or_path, repo_type, access_token) 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/rag/pipeline.py b/api/rag/pipeline.py new file mode 100644 index 000000000..e65758816 --- /dev/null +++ b/api/rag/pipeline.py @@ -0,0 +1,720 @@ +import logging +import os +import subprocess +from pathlib import 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 adalflow.core.types import Document, List +from adalflow.utils import get_adalflow_default_root_path + +from api.config import configs +from api.tools.embedder import get_embedder + +# Configure logging +logger = logging.getLogger(__name__) + +# 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: + """ + Count the number of tokens in a text string using tiktoken. + + Args: + text (str): The text to count tokens for. + embedder_type (str, optional): The embedder type ('openai', 'google', 'ollama', 'bedrock'). + If None, will be determined from configuration. + is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. + If None, will be determined from configuration. + + Returns: + int: The number of tokens in the text. + """ + 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 + + # 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": + # Ollama typically uses cl100k_base encoding + encoding = tiktoken.get_encoding("cl100k_base") + elif embedder_type == "google": + # Google uses similar tokenization to GPT models for rough estimation + encoding = tiktoken.get_encoding("cl100k_base") + 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 + # Use OpenAI embedding model encoding + encoding = tiktoken.encoding_for_model("text-embedding-3-small") + + return len(encoding.encode(text)) + except Exception as e: + # Fallback to a simple approximation if tiktoken fails + logger.warning(f"Error counting tokens with tiktoken: {e}") + # 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: + """ + Downloads a Git repository (GitHub, GitLab, or Bitbucket) to a specified local path. + + Args: + repo_type(str): Type of repository + repo_url (str): The URL of the Git repository to clone. + local_path (str): The local directory where the repository will be cloned. + access_token (str, optional): Access token for private repositories. + + Returns: + str: The output message from the `git` command. + """ + try: + # Check if Git is installed + logger.info(f"Preparing to clone repository to {local_path}") + subprocess.run( + ["git", "--version"], + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + # 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." + ) + return f"Using existing repository at {local_path}" + + # Ensure the local path exists + os.makedirs(local_path, exist_ok=True) + + # Prepare the clone URL with access token if provided + clone_url = repo_url + if access_token: + parsed = urlparse(repo_url) + # URL-encode the token to handle special characters + 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, + "", + "", + "", + ) + ) + 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, + "", + "", + "", + ) + ) + elif repo_type == "bitbucket": + # Bitbucket has two token formats with different auth schemes: + # - HTTP access tokens (prefix "ATCTT") use x-bitbucket-api-token-auth + # - App passwords (deprecated, EOL June 2026) use x-token-auth + # Detect by token prefix so existing app password users keep working. + if access_token.startswith("ATCTT"): + auth_scheme = "x-bitbucket-api-token-auth" + 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, + "", + "", + "", + ) + ) + + logger.info("Using access token for authentication") + + # Clone the repository + logger.info(f"Cloning repository from {repo_url} to {local_path}") + # We use repo_url in the log to avoid exposing the token in logs + result = subprocess.run( + ["git", "clone", "--depth=1", "--single-branch", clone_url, local_path], + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + logger.info("Repository cloned successfully") + return result.stdout.decode("utf-8") + + except subprocess.CalledProcessError as e: + 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="") + 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)}") + + +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. + + Args: + path (str): The root directory path. + embedder_type (str, optional): The embedder type ('openai', 'google', 'ollama'). + If None, will be determined from configuration. + is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. + If None, will be determined from configuration. + excluded_dirs (List[str], optional): List of directories to exclude from processing. + Overrides the default configuration if provided. + excluded_files (List[str], optional): List of file patterns to exclude from processing. + Overrides the default configuration if provided. + included_dirs (List[str], optional): List of directories to include exclusively. + When provided, only files in these directories will be processed. + included_files (List[str], optional): List of file patterns to include exclusively. + When provided, only files matching these patterns will be processed. + + Returns: + list: A list of Document objects with metadata. + """ + # Handle backward compatibility + if embedder_type is None and is_ollama_embedder is not None: + embedder_type = "ollama" if is_ollama_embedder else None + documents = [] + # File extensions to look for, prioritizing code files + 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) + + if use_inclusion_mode: + # Inclusion mode: only process specified directories and files + included_dirs = list(set(included_dirs)) if included_dirs else list() + included_files = list(set(included_files)) if included_files else list() + + logger.info("Using inclusion mode") + logger.info(f"Included directories: {included_dirs}") + logger.info(f"Included files: {included_files}") + + # Convert to lists for processing + excluded_dirs = [] + excluded_files = [] + else: + # Exclusion mode: use default exclusions plus any additional ones + 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: + final_excluded_dirs.update(excluded_dirs) + + if excluded_files is not None: + final_excluded_files.update(excluded_files) + + # Convert back to lists for compatibility + excluded_dirs = list(final_excluded_dirs) + excluded_files = list(final_excluded_files) + included_dirs = [] + included_files = [] + + 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}") + + for file_path in filter( + lambda p: _should_process_file( + p, + use_inclusion=use_inclusion_mode, + included_dirs=included_dirs, + included_files=included_files, + excluded_dirs=excluded_dirs, + excluded_files=excluded_files, + ), + 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, + }, + ) + documents.append(doc) + except Exception as e: + logger.error(f"Error reading {file_path}: {e}") + + logger.info(f"Found {len(documents)} documents") + return documents + +def prepare_data_pipeline(embedder_type: str = None, is_ollama_embedder: bool = None): + """ + Creates and returns the data transformation pipeline. + + Args: + embedder_type (str, optional): The embedder type ('openai', 'google', 'ollama'). + If None, will be determined from configuration. + is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. + If None, will be determined from configuration. + + Returns: + adal.Sequential: The data transformation pipeline + """ + from api.config import get_embedder_config, get_embedder_type + + # Handle backward compatibility + if embedder_type is None and is_ollama_embedder is not 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() + + splitter = TextSplitter(**configs["text_splitter"]) + embedder_config = get_embedder_config() + + embedder = get_embedder(embedder_type=embedder_type) + + batch_size = embedder_config.get("batch_size", 500) + 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, +) -> LocalDB: + """ + Transforms a list of documents and saves them to a local database. + + Args: + documents (list): A list of `Document` objects. + db_path (str): The path to the local database file. + embedder_type (str, optional): The embedder type ('openai', 'google', 'ollama'). + If None, will be determined from configuration. + is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. + If None, will be determined from configuration. + """ + # Get the data transformer + data_transformer = prepare_data_pipeline(embedder_type, is_ollama_embedder) + + # Save the documents to a local database + db = LocalDB() + db.register_transformer(transformer=data_transformer, key="split_and_embed") + db.load(documents) + db.transform(key="split_and_embed") + os.makedirs(os.path.dirname(db_path), exist_ok=True) + db.save_state(filepath=db_path) + return db + + +class DatabaseManager: + """ + Manages the creation, loading, transformation, and persistence of LocalDB instances. + """ + + def __init__(self): + self.db = None + 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]: + """ + Create a new database from the repository. + + Args: + repo_type(str): Type of repository + repo_url_or_path (str): The URL or local path of the repository + access_token (str, optional): Access token for private repositories + embedder_type (str, optional): Embedder type to use ('openai', 'google', 'ollama'). + If None, will be determined from configuration. + is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. + If None, will be determined from configuration. + excluded_dirs (List[str], optional): List of directories to exclude from processing + excluded_files (List[str], optional): List of file patterns to exclude from processing + included_dirs (List[str], optional): List of directories to include exclusively + included_files (List[str], optional): List of file patterns to include exclusively + + Returns: + List[Document]: List of Document objects + """ + # Handle backward compatibility + if embedder_type is None and is_ollama_embedder is not 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, + ) + + def reset_database(self): + """ + Reset the database to its initial state. + """ + self.db = None + self.repo_url_or_path = None + self.repo_paths = None + + 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("/") + + if repo_type in ["github", "gitlab", "bitbucket"] and len(url_parts) >= 5: + # GitHub URL format: https://github.com/owner/repo + # GitLab URL format: https://gitlab.com/owner/repo or https://gitlab.com/group/subgroup/repo + # Bitbucket URL format: https://bitbucket.org/owner/repo + owner = url_parts[-2] + repo = url_parts[-1].replace(".git", "") + repo_name = f"{owner}_{repo}" + else: + 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: + """ + Download and prepare all paths. + Paths: + ~/.adalflow/repos/{owner}_{repo_name} (for url, local path will be the same) + ~/.adalflow/databases/{owner}_{repo_name}.pkl + + Args: + repo_type(str): Type of repository + repo_url_or_path (str): The URL or local path of the repository + access_token (str, optional): Access token for private repositories + """ + logger.info(f"Preparing repo storage for {repo_url_or_path}...") + + 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://" + ): + # Extract the repository name from the URL + 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) + + # 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 + ) + else: + 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 + + save_db_file = os.path.join(root_path, "databases", f"{repo_name}.pkl") + os.makedirs(save_repo_dir, exist_ok=True) + os.makedirs(os.path.dirname(save_db_file), exist_ok=True) + + self.repo_paths = { + "save_repo_dir": save_repo_dir, + "save_db_file": save_db_file, + } + self.repo_url_or_path = repo_url_or_path + logger.info(f"Repo paths: {self.repo_paths}") + + except Exception as e: + 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]: + """ + Prepare the indexed database for the repository. + + Args: + embedder_type (str, optional): Embedder type to use ('openai', 'google', 'ollama'). + If None, will be determined from configuration. + is_ollama_embedder (bool, optional): DEPRECATED. Use embedder_type instead. + If None, will be determined from configuration. + excluded_dirs (List[str], optional): List of directories to exclude from processing + excluded_files (List[str], optional): List of file patterns to exclude from processing + included_dirs (List[str], optional): List of directories to include exclusively + included_files (List[str], optional): List of file patterns to include exclusively + + Returns: + List[Document]: List of Document objects + """ + + def _embedding_vector_length(doc: Document) -> int: + vector = getattr(doc, "vector", None) + if vector is None: + return 0 + try: + if hasattr(vector, "shape"): + if len(vector.shape) == 0: + return 0 + return int(vector.shape[-1]) + if hasattr(vector, "__len__"): + return int(len(vector)) + except Exception: + return 0 + return 0 + + # Handle backward compatibility + if embedder_type is None and is_ollama_embedder is not 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...") + try: + self.db = LocalDB.load_state(self.repo_paths["save_db_file"]) + documents = self.db.get_transformed_data(key="split_and_embed") + if documents: + lengths = [_embedding_vector_length(doc) for doc in documents] + non_empty = sum(1 for n in lengths if n > 0) + empty = len(lengths) - non_empty + sample_sizes = sorted({n for n in lengths if n > 0})[:3] + logger.info( + "Loaded %s documents from existing database (embeddings: %s non-empty, %s empty; sample_dims=%s)", + len(documents), + non_empty, + empty, + sample_sizes, + ) + + if non_empty == 0: + logger.warning( + "Existing database contains no usable embeddings. Rebuilding embeddings..." + ) + else: + return documents + except Exception as e: + logger.error(f"Error loading existing database: {e}") + # Continue to create a new database + + # prepare the database + logger.info("Creating new database...") + documents = read_all_documents( + self.repo_paths["save_repo_dir"], + embedder_type=embedder_type, + excluded_dirs=excluded_dirs, + excluded_files=excluded_files, + included_dirs=included_dirs, + included_files=included_files, + ) + self.db = transform_documents_and_save_to_db( + documents, self.repo_paths["save_db_file"], embedder_type=embedder_type + ) + logger.info(f"Total documents: {len(documents)}") + transformed_docs = self.db.get_transformed_data(key="split_and_embed") + 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 + ): + """ + Prepare the retriever for a repository. + This is a compatibility method for the isolated API. + + Args: + repo_type(str): Type of repository + repo_url_or_path (str): The URL or local path of the repository + access_token (str, optional): Access token for private repositories + + Returns: + List[Document]: List of Document objects + """ + return self.prepare_database(repo_url_or_path, repo_type, access_token) diff --git a/api/rag/rag.py b/api/rag/rag.py new file mode 100644 index 000000000..97fe4e91c --- /dev/null +++ b/api/rag/rag.py @@ -0,0 +1,390 @@ +import asyncio +import logging +import os +from collections import defaultdict +from collections.abc import Sized +from uuid import uuid4 + +import adalflow as adal +from adalflow.components.retriever.faiss_retriever import FAISSRetriever +from adalflow.core.types import ( + AssistantResponse, + DialogTurn, + Document, + RetrieverOutput, + 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 = { + 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), + ) + + +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 | list[str], language: str = "en" + ) -> list[RetrieverOutput]: + """ + Process a query using RAG. + + Args: + query: The user's query + + Returns: + list of RetrieverOutput. + """ + 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, returning empty list") + return [] + + 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) 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, 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)