diff --git a/.github/workflows/api_tests/mock_vllm.py b/.github/workflows/api_tests/mock_vllm.py index 0b057cd5b..724831afd 100644 --- a/.github/workflows/api_tests/mock_vllm.py +++ b/.github/workflows/api_tests/mock_vllm.py @@ -6,7 +6,7 @@ import hashlib import time import uuid -from typing import Any, List, Optional, Union +from typing import Any from fastapi import FastAPI from pydantic import BaseModel @@ -22,19 +22,19 @@ class EmbeddingRequest(BaseModel): model: str - input: Union[str, List[str]] + input: str | list[str] encoding_format: str = "float" class EmbeddingData(BaseModel): object: str = "embedding" - embedding: List[float] + embedding: list[float] index: int class EmbeddingResponse(BaseModel): object: str = "list" - data: List[EmbeddingData] + data: list[EmbeddingData] model: str usage: dict @@ -49,13 +49,13 @@ class ChatMessage(BaseModel): class ChatCompletionRequest(BaseModel): model: str - messages: List[ChatMessage] - temperature: Optional[float] = 0.7 - max_tokens: Optional[int] = 1024 - stream: Optional[bool] = False - top_p: Optional[float] = 1.0 - n: Optional[int] = 1 - stop: Optional[Union[str, List[str]]] = None + messages: list[ChatMessage] + temperature: float | None = 0.7 + max_tokens: int | None = 1024 + stream: bool | None = False + top_p: float | None = 1.0 + n: int | None = 1 + stop: str | list[str] | None = None class ChatCompletionChoice(BaseModel): @@ -75,7 +75,7 @@ class ChatCompletionResponse(BaseModel): object: str = "chat.completion" created: int model: str - choices: List[ChatCompletionChoice] + choices: list[ChatCompletionChoice] usage: ChatCompletionUsage @@ -84,13 +84,13 @@ class ChatCompletionResponse(BaseModel): class TextCompletionRequest(BaseModel): model: str - prompt: Union[str, List[str]] - temperature: Optional[float] = 0.7 - max_tokens: Optional[int] = 1024 - stream: Optional[bool] = False - top_p: Optional[float] = 1.0 - n: Optional[int] = 1 - stop: Optional[Union[str, List[str]]] = None + prompt: str | list[str] + temperature: float | None = 0.7 + max_tokens: int | None = 1024 + stream: bool | None = False + top_p: float | None = 1.0 + n: int | None = 1 + stop: str | list[str] | None = None class TextCompletionChoice(BaseModel): @@ -104,14 +104,14 @@ class TextCompletionResponse(BaseModel): object: str = "text_completion" created: int model: str - choices: List[TextCompletionChoice] + choices: list[TextCompletionChoice] usage: ChatCompletionUsage # Same structure as chat # ============== Helper Functions ============== -def generate_fake_embedding(text: str, dim: int = EMBEDDING_DIM) -> List[float]: +def generate_fake_embedding(text: str, dim: int = EMBEDDING_DIM) -> list[float]: """Generate deterministic fake embedding based on text hash.""" h = hashlib.md5(text.encode()).digest() result = [] @@ -128,7 +128,7 @@ def count_tokens(text: str) -> int: return 10 # Default for non-string content -def generate_mock_response(messages: List[ChatMessage]) -> str: +def generate_mock_response(messages: list[ChatMessage]) -> str: """Generate a mock response based on the input messages.""" last_message = messages[-1] if messages else None if not last_message: @@ -180,10 +180,7 @@ async def list_models(): async def create_embeddings(request: EmbeddingRequest) -> EmbeddingResponse: inputs = request.input if isinstance(request.input, list) else [request.input] - data = [ - EmbeddingData(embedding=generate_fake_embedding(text), index=i) - for i, text in enumerate(inputs) - ] + data = [EmbeddingData(embedding=generate_fake_embedding(text), index=i) for i, text in enumerate(inputs)] return EmbeddingResponse( data=data, @@ -196,9 +193,7 @@ async def create_embeddings(request: EmbeddingRequest) -> EmbeddingResponse: async def create_chat_completion(request: ChatCompletionRequest) -> ChatCompletionResponse: """Mock chat completion endpoint for LLM/VLM requests.""" # Calculate token counts - prompt_tokens = sum( - count_tokens(str(msg.content)) for msg in request.messages - ) + prompt_tokens = sum(count_tokens(str(msg.content)) for msg in request.messages) # Generate mock response response_text = generate_mock_response(request.messages) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 000000000..aef59962c --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,52 @@ +name: Linting + +on: + push: + branches: [ main, dev ] + + # Temporarily disabled until merge + # pull_request: + # branches: [ main, dev ] + +jobs: + lint: + runs-on: ubuntu-latest + + strategy: + matrix: + python-version: ["3.12"] + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install uv + run: | + curl -LsSf https://astral.sh/uv/install.sh | sh + echo "$HOME/.local/bin" >> $GITHUB_PATH + + - name: Cache uv + uses: actions/cache@v4 + with: + path: | + ~/.cache/uv + key: uv-${{ runner.os }}-${{ matrix.python-version }}-${{ hashFiles('**/pyproject.toml', '**/uv.lock') }} + restore-keys: | + uv-${{ runner.os }}-${{ matrix.python-version }}- + + - name: Install dependencies with uv + run: | + uv sync --group lint + + - name: Run Ruff linting + run: | + uv run ruff check . --output-format=github + + - name: Check formatting + run: | + uv run ruff format --check . diff --git a/LINTING.md b/LINTING.md new file mode 100644 index 000000000..cd0b01af0 --- /dev/null +++ b/LINTING.md @@ -0,0 +1,148 @@ +# Linting Setup Guide + +This project uses **Ruff** for Python linting and code formatting. All developers must ensure their code passes linting checks before submitting pull requests. + +## Installation + +The linting dependencies are already included in `pyproject.toml`. Install them with: + +```bash +uv sync --group lint +``` + +Or install all groups including dev: + +```bash +uv sync +``` + +## Running Linting Locally + +### Check for linting issues + +```bash +uv run ruff check . +``` + +### Auto-fix linting issues + +```bash +uv run ruff check . --fix +``` + +### Check code formatting + +```bash +uv run ruff format --check . +``` + +### Auto-format code + +```bash +uv run ruff format . +``` + +## Ruff Configuration + +Ruff configuration is defined in `pyproject.toml` under `[tool.ruff]` and `[tool.ruff.lint]` sections. + +### Key Rules Enabled + +- **E, W**: PEP 8 errors and warnings +- **F**: Pyflakes (undefined names, unused imports) +- **I**: isort-style import sorting +- **C4**: flake8-comprehensions (optimization suggestions) +- **UP**: pyupgrade (modernize Python code) +- **PIE**: flake8-pie (misc linting rules) +- **RUF**: Ruff-specific rules (**`not applied currently`**) + +### Excluded Directories + +The following directories are excluded from linting: +- `.git`, `.venv`, `.pytest_cache`, `.ruff_cache`, `.vscode` +- `node_modules`, `vdb`, `ray_mount`, `model_weights`, `logs` +- `.astro`, `openrag.egg-info` + +## CI/CD Integration + +A GitHub Actions workflow (`.github/workflows/lint.yml`) automatically runs on every push and pull request. The workflow: + +1. Checks for linting issues with `ruff check` +2. Verifies code formatting with `ruff format --check` +3. Reports results directly to GitHub (with file and line annotations) + +### Pull Request Requirements + +All pull requests must pass the linting CI check. If your PR fails: + +1. Run `uv run ruff check . --fix` to auto-fix issues +2. Run `uv run ruff format .` to format the code +3. Review any remaining issues that require manual fixes +4. Commit and push your changes + +## Pre-commit Hook (Optional) +One can also setup a pre-commit hook to run linting automatically before each commit. + +```bash +cat > .git/hooks/pre-commit << 'EOF' +#!/bin/bash + +# Get list of staged Python files +FILES=$(git diff --cached --name-only --diff-filter=ACM | grep '\.py$') + +if [ -z "$FILES" ]; then + exit 0 +fi + +# Run ruff check (will auto-fix if possible) +echo "Running ruff check..." +uv run ruff check $FILES --fix + +# Run ruff format +echo "Running ruff format..." +uv run ruff format $FILES + +# Re-add the modified files +git add $FILES + +# Run ruff check again without --fix to catch remaining errors +echo "Checking for remaining issues..." +uv run ruff check $FILES + +# If ruff check fails, block the commit +if [ $? -ne 0 ]; then + echo "Ruff check failed. Please fix errors before committing." + exit 1 +fi + +exit 0 +EOF + +chmod +x .git/hooks/pre-commit +``` + +* To disable pre-commit + +```bash +rm .git/hooks/pre-commit +``` + +## Troubleshooting + +### "ruff: command not found" +Make sure you've installed the lint dependencies: +```bash +uv sync --group lint +``` + +### Auto-fix doesn't resolve all issues +Some issues require manual fixes. Review the error messages and fix them manually, then commit again. +Sometimes, you will have to ignore a particular check if deemed unnecessary + +```python +# Example: C417 rule suggest that map here is unnecessary + + +docs_with_tokens = list(map(lambda d: (_length_function(d.page_content), d), docs)) # noqa: C417 + +``` \ No newline at end of file diff --git a/automatic-evaluation-pipeline/benchmark.py b/automatic-evaluation-pipeline/benchmark.py index 9733339c1..1c5d943c8 100644 --- a/automatic-evaluation-pipeline/benchmark.py +++ b/automatic-evaluation-pipeline/benchmark.py @@ -1,16 +1,17 @@ -import os import asyncio import json import math -from typing import Literal, TypedDict -from loguru import logger +import os +from typing import TypedDict + import numpy as np -from openai import AsyncOpenAI +import pandas as pd +from dotenv import load_dotenv from langchain_openai import ChatOpenAI +from loguru import logger +from openai import AsyncOpenAI from pydantic import BaseModel, Field -from dotenv import load_dotenv from tqdm.asyncio import tqdm -import pandas as pd load_dotenv() @@ -49,6 +50,7 @@ async def retrieve_response_and_docs_openrag( logger.debug(f"Error fetching chunks and response: {e}") return None, [] + def compute_hits(true_chunk_id, all_retrieved_chunks): return true_chunk_id in all_retrieved_chunks @@ -67,6 +69,7 @@ def compute_inverted_ranks(true_chunk_id, all_retrieved_chunks): else: return 0 + # Sources retrieval evaluation def relevance(val, true_chunk_ids): return 1 if val in true_chunk_ids else 0 @@ -99,11 +102,12 @@ def source_score_per_question( class CompletionEvaluationResponse(BaseModel): - score: int = Field(..., - ge=1, - le=10, - description="Le résultat du juge LLM." - "Le résultat est compris entre 1 et 10, où 1 indique une réponse très incomplète et 10 indique une réponse très complète." + score: int = Field( + ..., + ge=1, + le=10, + description="Le résultat du juge LLM." + "Le résultat est compris entre 1 et 10, où 1 indique une réponse très incomplète et 10 indique une réponse très complète.", # "The output of the LLM judge. It can be one of the following: " # "'complete', 'mostly_complete', 'partially_complete', 'incomplete'. " # "This indicates how well the generated answer matches the true answer.", @@ -111,11 +115,12 @@ class CompletionEvaluationResponse(BaseModel): class PrecisionEvaluationResponse(BaseModel): - score: int = Field(..., - ge=1, - le=10, - description="Le résultat du juge LLM." - "Le résultat est compris entre 1 et 10, où 1 indique une réponse très imprécise et 10 indique une réponse très précise." + score: int = Field( + ..., + ge=1, + le=10, + description="Le résultat du juge LLM." + "Le résultat est compris entre 1 et 10, où 1 indique une réponse très imprécise et 10 indique une réponse très précise.", # "The output of the LLM judge. It can be one of the following: " # "'Highly_precise', 'mostly_precise', 'low_precision', 'imprecise'. " # "This indicates how well the generated answer matches the true answer.", @@ -214,7 +219,7 @@ class Element(TypedDict): async def main(): - with open("./dataset.json", "r", encoding="utf-8") as f: + with open("./dataset.json", encoding="utf-8") as f: eval_dataset: list[Element] = json.load(f) list_response_answer_reference = eval_dataset # [:10] @@ -222,7 +227,7 @@ async def main(): num_port = os.environ.get("APP_PORT") num_host = os.environ["APP_URL"] openrag_api_base_url = f"http://{num_host}:{num_port}" - partition = "pdftest" # To replace with your wanted partition's name + partition = "pdftest" # To replace with your wanted partition's name # Create shared semaphores for rate limiting openrag_semaphore = asyncio.Semaphore(4) # Limit concurrent OpenRAG requests @@ -248,14 +253,18 @@ async def main(): ): if openrag_response is None: continue - chunk_id_reference = [c["id"] for c in input_reference["chunks"]] # The "true answer" ids list + chunk_id_reference = [ + c["id"] for c in input_reference["chunks"] + ] # The "true answer" ids list # Hit rate and MRR hit_rates.append(compute_hits(chunk_id_reference[0], openrag_chunk_ids)) MRRs.append(compute_inverted_ranks(chunk_id_reference[0], openrag_chunk_ids)) # Recall computaton - recall = len(list(set(chunk_id_reference) & set(openrag_chunk_ids))) / len(chunk_id_reference) + recall = len(list(set(chunk_id_reference) & set(openrag_chunk_ids))) / len( + chunk_id_reference + ) recalls.append(recall) # nDCG score calculation @@ -281,7 +290,7 @@ async def main(): print(f"Average Hit Rate: {round(np.array(hit_rates).mean(), 3)}") print(f"Average MRR: {round(np.array(MRRs).mean(), 3)}") print(f"Average Recall: {round(np.array(recalls).mean(), 3)}") - + # Filter out error responses valid_scores = [(comp, prec) for comp, prec in llm_judge_scores if comp != "error"] valid_ndcg_scores = nDCG_scores[: len(valid_scores)] # Match the filtered scores @@ -310,11 +319,16 @@ async def main(): print("\n", "-" * 50, "\n") print("\nCompletion evaluation distribution:") print(eval_results["completion_evaluation"].value_counts()) - print(f"Completion evaluation average: {eval_results['completion_evaluation'].mean():.3f}") + print( + f"Completion evaluation average: {eval_results['completion_evaluation'].mean():.3f}" + ) print("\n", "-" * 50, "\n") print("\nPrecision evaluation distribution:") print(eval_results["precision_evaluation"].value_counts()) - print(f"Precision evaluation average: {eval_results['precision_evaluation'].mean():.3f}") + print( + f"Precision evaluation average: {eval_results['precision_evaluation'].mean():.3f}" + ) + if __name__ == "__main__": asyncio.run(main()) diff --git a/automatic-evaluation-pipeline/generate_questions.py b/automatic-evaluation-pipeline/generate_questions.py index d85259307..a8edb46ac 100644 --- a/automatic-evaluation-pipeline/generate_questions.py +++ b/automatic-evaluation-pipeline/generate_questions.py @@ -1,20 +1,18 @@ -import random +import ast import asyncio -import os -import httpx import json +import os +import random import time -import numpy as np -import ast +import hdbscan +import httpx +import numpy as np from dotenv import load_dotenv -from loguru import logger from langchain_openai import ChatOpenAI -from tqdm.asyncio import tqdm +from loguru import logger from sklearn.cluster import DBSCAN -import umap.umap_ as umap -import hdbscan -import pickle +from tqdm.asyncio import tqdm load_dotenv() # Charge les variables du .env @@ -63,7 +61,7 @@ async def summarize( message = [ { "role": "user", - "content": f"Voici le document:\n{chunk}. Donnez-moi un résumé qui précise quel type d'informations et de contenu contient le passage, mais sans entrer dans des détails trop précis." + "content": f"Voici le document:\n{chunk}. Donnez-moi un résumé qui précise quel type d'informations et de contenu contient le passage, mais sans entrer dans des détails trop précis.", } ] @@ -149,7 +147,7 @@ async def main(): num_port = os.environ.get("APP_PORT") num_host = os.environ["APP_URL"] openrag_api_base_url = f"http://{num_host}:{num_port}" - partition = "pdftest" # To replace with your wanted partition's name + partition = "pdftest" # To replace with your wanted partition's name url = f"{openrag_api_base_url}/partition/{partition}/chunks" start = time.time() diff --git a/automatic-evaluation-pipeline/upload_files.py b/automatic-evaluation-pipeline/upload_files.py index 04baf34a0..d61afec7e 100644 --- a/automatic-evaluation-pipeline/upload_files.py +++ b/automatic-evaluation-pipeline/upload_files.py @@ -1,8 +1,9 @@ -import httpx import os -from loguru import logger from pathlib import Path + +import httpx from dotenv import load_dotenv +from loguru import logger load_dotenv() @@ -13,6 +14,7 @@ dir_name = "./pdf_files" # Replace with your directory path dir_path = Path(dir_name).resolve() + def __check_api(base_url): try: response = httpx.get(f"{base_url}/health_check") diff --git a/openrag/api.py b/openrag/api.py index a3f52821b..349069bfb 100644 --- a/openrag/api.py +++ b/openrag/api.py @@ -1,37 +1,25 @@ import os import warnings - -# Filter SyntaxWarning from pydub (invalid escape sequences in regex) -# This is a known issue in pydub 0.25.1 that hasn't been fixed upstream -warnings.filterwarnings("ignore", category=SyntaxWarning, module="pydub") - -import ray -from dotenv import dotenv_values - -SHARED_ENV = os.environ.get("SHARED_ENV", None) - -env_vars = dotenv_values(SHARED_ENV) if SHARED_ENV else {} -env_vars["PYTHONPATH"] = "/app/openrag" - - -ray.init(dashboard_host="0.0.0.0") - - -import os from enum import Enum +from importlib.metadata import version as get_package_version from pathlib import Path -from typing import Optional +import ray import uvicorn -from importlib.metadata import version as get_package_version - from config import load_config - +from dotenv import dotenv_values from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.openapi.utils import get_openapi from fastapi.responses import JSONResponse from fastapi.staticfiles import StaticFiles + +ray.init(dashboard_host="0.0.0.0") + +# Apply noqa: E402 to ignore "module level import not at top of file" cause ray.init has to be called first + +# flake8: noqa: E402 + from routers.actors import router as actors_router from routers.extract import router as extract_router from routers.indexer import router as indexer_router @@ -39,13 +27,24 @@ from routers.partition import router as partition_router from routers.queue import router as queue_router from routers.search import router as search_router -from routers.users import router as users_router from routers.tools import router as tools_router +from routers.users import router as users_router from starlette.middleware.base import BaseHTTPMiddleware from utils.dependencies import get_vectordb from utils.exceptions import OpenRAGError from utils.logger import get_logger +# Filter SyntaxWarning from pydub (invalid escape sequences in regex) +# This is a known issue in pydub 0.25.1 that hasn't been fixed upstream +warnings.filterwarnings("ignore", category=SyntaxWarning, module="pydub") + + +SHARED_ENV = os.environ.get("SHARED_ENV", None) + +env_vars = dotenv_values(SHARED_ENV) if SHARED_ENV else {} +env_vars["PYTHONPATH"] = "/app/openrag" + + logger = get_logger() config = load_config() DATA_DIR = Path(config.paths.data_dir) @@ -71,15 +70,12 @@ def __init__(self, config): # Read the token from env (or None if not set) -AUTH_TOKEN: Optional[str] = os.getenv("AUTH_TOKEN") -INDEXERUI_PORT: Optional[str] = os.getenv("INDEXERUI_PORT", "3042") -INDEXERUI_URL: Optional[str] = os.getenv( - "INDEXERUI_URL", f"http://localhost:{INDEXERUI_PORT}" -) -WITH_CHAINLIT_UI: Optional[bool] = ( - os.getenv("WITH_CHAINLIT_UI", "true").lower() == "true" -) -WITH_OPENAI_API: Optional[bool] = os.getenv("WITH_OPENAI_API", "true").lower() == "true" +AUTH_TOKEN: str | None = os.getenv("AUTH_TOKEN") +INDEXERUI_PORT: str | None = os.getenv("INDEXERUI_PORT", "3042") +INDEXERUI_URL: str | None = os.getenv("INDEXERUI_URL", f"http://localhost:{INDEXERUI_PORT}") +WITH_CHAINLIT_UI: bool = os.getenv("WITH_CHAINLIT_UI", "true").lower() == "true" +WITH_OPENAI_API: bool = os.getenv("WITH_OPENAI_API", "true").lower() == "true" + try: app_version = get_package_version("openrag") @@ -98,9 +94,7 @@ def custom_openapi(): routes=app.routes, ) # Add global security - openapi_schema["components"]["securitySchemes"] = { - "BearerAuth": {"type": "http", "scheme": "bearer"} - } + openapi_schema["components"]["securitySchemes"] = {"BearerAuth": {"type": "http", "scheme": "bearer"}} openapi_schema["security"] = [{"BearerAuth": []}] app.openapi_schema = openapi_schema return app.openapi_schema @@ -126,7 +120,7 @@ async def dispatch(self, request: Request, call_next): "/openapi.json", "/redoc", "/health_check", - "/version" + "/version", ] or request.url.path.startswith("/chainlit"): # Allow all chainlit subroutes return await call_next(request) @@ -189,22 +183,16 @@ async def openrag_exception_handler(request: Request, exc: OpenRAGError): ) app.state.app_state = AppState(config) -app.mount( - "/static", StaticFiles(directory=DATA_DIR.resolve(), check_dir=True), name="static" -) +app.mount("/static", StaticFiles(directory=DATA_DIR.resolve(), check_dir=True), name="static") -@app.get( - "/health_check", summary="Health check endpoint for API", dependencies=[] -) +@app.get("/health_check", summary="Health check endpoint for API", dependencies=[]) async def health_check(request: Request): # TODO : Error reporting about llm and vlm return "RAG API is up." -@app.get( - "/version", summary="Get openRAG version", dependencies=[] -) +@app.get("/version", summary="Get openRAG version", dependencies=[]) def get_version(): return {"version": app.version} @@ -235,9 +223,7 @@ def get_version(): from chainlit.utils import mount_chainlit mount_chainlit(app, "./app_front.py", path="/chainlit") - app.include_router( - openai_router, prefix="/v1", tags=[Tags.OPENAI] - ) # cause chainlit uses openai api endpoints + app.include_router(openai_router, prefix="/v1", tags=[Tags.OPENAI]) # cause chainlit uses openai api endpoints if __name__ == "__main__": if config.ray.serve.enable: @@ -248,21 +234,14 @@ def get_version(): class OpenRagAPI: pass - serve.start( - http_options={"host": config.ray.serve.host, "port": config.ray.serve.port} - ) + serve.start(http_options={"host": config.ray.serve.host, "port": config.ray.serve.port}) if WITH_CHAINLIT_UI: from chainlit_api import app as chainlit_app + serve.run(OpenRagAPI.bind(), route_prefix="/") - uvicorn.run( - chainlit_app, - host="0.0.0.0", - port=config.ray.serve.chainlit_port - ) + uvicorn.run(chainlit_app, host="0.0.0.0", port=config.ray.serve.chainlit_port) else: serve.run(OpenRagAPI.bind(), route_prefix="/", blocking=True) else: - uvicorn.run( - "api:app", host="0.0.0.0", port=8080, reload=True, proxy_headers=True - ) + uvicorn.run("api:app", host="0.0.0.0", port=8080, reload=True, proxy_headers=True) diff --git a/openrag/app_front.py b/openrag/app_front.py index 262df05d2..9de2cc2be 100644 --- a/openrag/app_front.py +++ b/openrag/app_front.py @@ -129,7 +129,7 @@ async def chat_profile(current_user: cl.User): ) return chat_profiles except Exception as e: - await cl.Message(content=f"An error occured: {str(e)}").send() + await cl.Message(content=f"An error occured: {e!s}").send() @cl.on_chat_start @@ -151,7 +151,7 @@ async def on_chat_start(): except Exception as e: logger.exception("An error occured while checking the API health", error=str(e)) await cl.Message( - content=f"An error occured while checking the API health: {str(e)}" + content=f"An error occured while checking the API health: {e!s}" ).send() @@ -281,4 +281,4 @@ async def on_message(message: cl.Message): await msg.update() except Exception as e: logger.exception("Error during chat completion", error=str(e)) - await cl.Message(content=f"An error occurred: {str(e)}").send() + await cl.Message(content=f"An error occurred: {e!s}").send() diff --git a/openrag/chainlit_api.py b/openrag/chainlit_api.py index a38ddf88f..58e8cde3f 100644 --- a/openrag/chainlit_api.py +++ b/openrag/chainlit_api.py @@ -1,6 +1,6 @@ -from fastapi import FastAPI from chainlit.utils import mount_chainlit +from fastapi import FastAPI app = FastAPI() -mount_chainlit(app=app, target="./app_front.py", path="/chainlit") \ No newline at end of file +mount_chainlit(app=app, target="./app_front.py", path="/chainlit") diff --git a/openrag/components/indexer/chunker/chunker.py b/openrag/components/indexer/chunker/chunker.py index 15b8f99ac..3a8e5f1d1 100644 --- a/openrag/components/indexer/chunker/chunker.py +++ b/openrag/components/indexer/chunker/chunker.py @@ -1,7 +1,6 @@ -from typing import Literal, Optional +from typing import Literal import openai - from components.indexer.utils.text_sanitizer import sanitize_text from components.prompts import CHUNK_CONTEXTUALIZER_PROMPT from components.utils import detect_language, get_vlm_semaphore, load_config @@ -23,13 +22,9 @@ # Timeout for individual chunk contextualization LLM calls (in seconds) CONTEXTUALIZATION_TIMEOUT = config.chunker.get("contextualization_timeout", 120) # Maximum concurrent contextualization tasks to prevent system overload -MAX_CONCURRENT_CONTEXTUALIZATION = config.chunker.get( - "max_concurrent_contextualization", 10 -) +MAX_CONCURRENT_CONTEXTUALIZATION = config.chunker.get("max_concurrent_contextualization", 10) -BASE_CHUNK_FORMAT = ( - "* filename: {filename}\n\n[CHUNK_START]\n\n{content}\n\n[CHUNK_END]" -) +BASE_CHUNK_FORMAT = "* filename: {filename}\n\n[CHUNK_START]\n\n{content}\n\n[CHUNK_END]" CHUNK_FORMAT = "[CONTEXT]\n\n{chunk_context}\n\n" + BASE_CHUNK_FORMAT @@ -54,10 +49,10 @@ async def _generate_context( user_msg = f""" Here is the context to consider for generating the context: - Filename: {filename} - - First chunks: + - First chunks: {"\n--\n".join(c.page_content for c in first_chunks)} - - Previous chunks: + - Previous chunks: {"\n--\n".join(c.page_content for c in prev_chunks)} Here is the current chunk to contextualize strictly in this {lang} language: @@ -119,8 +114,7 @@ async def contextualize_chunks( batch_contexts = await tqdm.gather( *batch_tasks, total=len(batch_tasks), - desc=f"Contextualizing chunks of *{filename}* " - f"[{batch_start + 1}-{batch_end}/{len(chunks)}]", + desc=f"Contextualizing chunks of *{filename}* [{batch_start + 1}-{batch_end}/{len(chunks)}]", ) contexts.extend(batch_contexts) @@ -148,7 +142,7 @@ def __init__( self, chunk_size: int = 200, chunk_overlap_rate: float = 0.2, - llm_config: Optional[dict] = None, + llm_config: dict | None = None, contextual_retrieval: bool = False, **kwargs, ): @@ -164,9 +158,7 @@ def __init__( self.contextual_retrieval = contextual_retrieval # Initialize contextualizer only if needed - self.contextualizer = ( - ChunkContextualizer(llm_config) if contextual_retrieval else None - ) + self.contextualizer = ChunkContextualizer(llm_config) if contextual_retrieval else None async def _apply_contextualization( self, @@ -178,21 +170,15 @@ async def _apply_contextualization( if not self.contextual_retrieval or len(chunks) < 2: return [ Document( - page_content=BASE_CHUNK_FORMAT.format( - chunk_context="", filename=filename, content=c.page_content - ), + page_content=BASE_CHUNK_FORMAT.format(chunk_context="", filename=filename, content=c.page_content), metadata=c.metadata, ) for c in chunks ] - return await self.contextualizer.contextualize_chunks( - chunks, lang=lang, filename=filename - ) + return await self.contextualizer.contextualize_chunks(chunks, lang=lang, filename=filename) - def _prepare_md_elements( - self, content: str - ) -> tuple[list[MDElement], list[MDElement]]: + def _prepare_md_elements(self, content: str) -> tuple[list[MDElement], list[MDElement]]: """Prepare and combine markdown elements from raw content.""" md_elements: list[MDElement] = split_md_elements(content) @@ -200,14 +186,10 @@ def _prepare_md_elements( for e in md_elements: if e.type in ("table", "image"): - if ( - e.type == "image" and IMAGE_PLACEHOLDER.lower() in e.content.lower() - ): # skip placeholder images + if e.type == "image" and IMAGE_PLACEHOLDER.lower() in e.content.lower(): # skip placeholder images continue - if ( - self._length_function(e.content) <= 100 - ): # do not isolate small tables/images + if self._length_function(e.content) <= 100: # do not isolate small tables/images texts.append(e) else: tables_and_images.append(e) @@ -219,9 +201,7 @@ def _prepare_md_elements( def split_text(self, text: str) -> list[str]: """Split text into chunks using the text splitter.""" if not self.text_splitter: - logger.warning( - "Text splitter not initialized. Initializing with default RecursiveCharacterTextSplitter." - ) + logger.warning("Text splitter not initialized. Initializing with default RecursiveCharacterTextSplitter.") from langchain.text_splitter import RecursiveCharacterTextSplitter self.text_splitter = RecursiveCharacterTextSplitter( @@ -232,9 +212,7 @@ def split_text(self, text: str) -> list[str]: return self.text_splitter.split_text(text) - def _get_chunks( - self, content: str, metadata: Optional[dict] = None, log=None - ) -> list[Document]: + def _get_chunks(self, content: str, metadata: dict | None = None, log=None) -> list[Document]: log = log or logger texts, tables_and_images = self._prepare_md_elements(content=content) combined_texts = "\n".join([e.content for e in texts]) @@ -290,9 +268,7 @@ def _get_chunks( prev_page_num = 1 for c in text_chunks: - page_info = get_chunk_page_number( - chunk_str=c, previous_chunk_ending_page=prev_page_num - ) + page_info = get_chunk_page_number(chunk_str=c, previous_chunk_ending_page=prev_page_num) start_page = page_info["start_page"] prev_page_num = page_info["end_page"] chunks.append( @@ -306,14 +282,10 @@ def _get_chunks( chunks.sort(key=lambda d: d.metadata.get("page")) return chunks else: - log.warning( - "No chunks created. Content is empty or image is not informative." - ) + log.warning("No chunks created. Content is empty or image is not informative.") return [] - async def split_document( - self, doc: Document, task_id: Optional[str] = None - ) -> list[Document]: + async def split_document(self, doc: Document, task_id: str | None = None) -> list[Document]: """Split document into chunks with optional contextualization.""" metadata = doc.metadata filename = metadata.get("filename", "") @@ -335,9 +307,7 @@ async def split_document( "Contextualizing chunks", apply_contextualization=self.contextual_retrieval, ) - chunks = await self._apply_contextualization( - chunks, lang=detected_lang, filename=filename - ) + chunks = await self._apply_contextualization(chunks, lang=detected_lang, filename=filename) log.info("Document chunking completed") return chunks else: @@ -353,9 +323,7 @@ def __init__( contextual_retrieval=False, **kwargs, ): - super().__init__( - chunk_size, chunk_overlap_rate, llm_config, contextual_retrieval, **kwargs - ) + super().__init__(chunk_size, chunk_overlap_rate, llm_config, contextual_retrieval, **kwargs) from langchain.text_splitter import RecursiveCharacterTextSplitter @@ -376,7 +344,7 @@ class ChunkerFactory: @staticmethod def create_chunker( config: OmegaConf, - embedder: Optional[BaseEmbedding] = None, + embedder: BaseEmbedding | None = None, ) -> BaseChunker: # Extract parameters chunker_params = OmegaConf.to_container(config.chunker, resolve=True) @@ -387,8 +355,7 @@ def create_chunker( if not chunker_cls: raise ValueError( - f"Chunker '{name}' is not recognized." - f" Available chunkers: {list(ChunkerFactory.CHUNKERS.keys())}" + f"Chunker '{name}' is not recognized. Available chunkers: {list(ChunkerFactory.CHUNKERS.keys())}" ) chunker_params["llm_config"] = config.vlm diff --git a/openrag/components/indexer/chunker/utils.py b/openrag/components/indexer/chunker/utils.py index 04bd2fd00..44f599c4f 100644 --- a/openrag/components/indexer/chunker/utils.py +++ b/openrag/components/indexer/chunker/utils.py @@ -1,5 +1,6 @@ import re -from typing import Callable, Literal, Optional +from collections.abc import Callable +from typing import Literal from components.indexer.utils.text_sanitizer import clean_markdown_table_spacing @@ -23,7 +24,7 @@ def __init__( self, type: Literal["text", "table", "image"], content: str, - page_number: Optional[int] = None, + page_number: int | None = None, ): self.type = type # 'text', 'table', 'image' self.content = content @@ -193,7 +194,7 @@ def parse_markdown_table(markdown_table): def chunk_table( table_element: MDElement, chunk_size: int = 512, - length_function: Optional[Callable[[str], int]] = None, + length_function: Callable[[str], int] | None = None, ) -> list[MDElement]: txt = clean_markdown_table_spacing(table_element.content) header_lines, groups = parse_markdown_table(txt) diff --git a/openrag/components/indexer/embeddings/__init__.py b/openrag/components/indexer/embeddings/__init__.py index fa3ff34d1..110056779 100644 --- a/openrag/components/indexer/embeddings/__init__.py +++ b/openrag/components/indexer/embeddings/__init__.py @@ -18,4 +18,4 @@ def get_embedder(embeddings_config: dict) -> BaseEmbedding: return embedder_class(embeddings_config) -__all__ = ["EmbeddingFactory", "BaseEmbedding", "OpenAIEmbedding"] +__all__ = ["BaseEmbedding", "EmbeddingFactory", "OpenAIEmbedding"] diff --git a/openrag/components/indexer/embeddings/openai.py b/openrag/components/indexer/embeddings/openai.py index 2a9dfde6e..9dc031e74 100644 --- a/openrag/components/indexer/embeddings/openai.py +++ b/openrag/components/indexer/embeddings/openai.py @@ -44,7 +44,7 @@ def embed_documents(self, texts: list[str | Document]) -> list[list[float]]: except openai.APIError as e: logger.error("API error in embed_documents", error=str(e)) raise EmbeddingAPIError( - f"OpenAI API error during document embedding: {str(e)}", + f"OpenAI API error during document embedding: {e!s}", model_name=self.embedding_model, base_url=self.base_url, error=str(e), @@ -62,7 +62,7 @@ def embed_documents(self, texts: list[str | Document]) -> list[list[float]]: except Exception as e: logger.exception("Unexpected error while embedding documents", error=str(e)) raise UnexpectedEmbeddingError( - f"Failed to embed documents: {str(e)}", + f"Failed to embed documents: {e!s}", model_name=self.embedding_model, base_url=self.base_url, error=str(e), diff --git a/openrag/components/indexer/indexer.py b/openrag/components/indexer/indexer.py index df4e337cb..97c495d3e 100644 --- a/openrag/components/indexer/indexer.py +++ b/openrag/components/indexer/indexer.py @@ -4,7 +4,7 @@ import traceback from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, List, Optional, Union +from typing import Any import ray import torch @@ -50,17 +50,15 @@ def __init__(self): self.logger.info("Indexer actor initialized.") @ray.method(concurrency_group="chunk") - async def chunk( - self, doc: Document, file_path: str, task_id: str = None - ) -> List[Document]: + async def chunk(self, doc: Document, file_path: str, task_id: str = None) -> list[Document]: chunks = await self.chunker.split_document(doc, task_id) return chunks @ray.method(concurrency_group="serialize") async def serialize_file( self, - path: Union[str, List[str]], - metadata: Optional[Dict] = {}, + path: str, + metadata: dict = {}, task_id: str = None, ): # Serialize @@ -69,21 +67,21 @@ async def serialize_file( async def add_file( self, - path: Union[str, List[str]], - metadata: Optional[Dict] = {}, - partition: Optional[str] = None, - user: Optional[Dict] = None, + path: str, + metadata: dict | None = None, + partition: str | None = None, + user: dict | None = None, ): task_state_manager = ray.get_actor("TaskStateManager", namespace="openrag") task_id = ray.get_runtime_context().get_task_id() + metadata = metadata or {} + file_id = metadata.get("file_id", None) log = self.logger.bind(file_id=file_id, partition=partition, task_id=task_id) log.info("Queued file for indexing.") try: # Set task details - user_metadata = { - k: v for k, v in metadata.items() if k not in {"file_id", "source"} - } + user_metadata = {k: v for k, v in metadata.items() if k not in {"file_id", "source"}} await task_state_manager.set_details.remote( task_id, @@ -98,18 +96,14 @@ async def add_file( metadata = {**metadata, "partition": partition} # Serialize - doc = await self.serialize_file( - path=path, metadata=metadata, task_id=task_id - ) + doc = await self.serialize_file(path=path, metadata=metadata, task_id=task_id) # Chunk if doc: await task_state_manager.set_state.remote(task_id, "CHUNKING") chunks = await self.handle.chunk.remote(doc, str(path), task_id) else: - log.warning( - "No document returned from serialization; skipping indexing." - ) + log.warning("No document returned from serialization; skipping indexing.") chunks = [] if self.enable_insertion: @@ -118,13 +112,9 @@ async def add_file( await self.handle.insert_documents.remote(chunks, user=user) log.info(f"Document {path} indexed successfully") else: - log.debug( - "No chunks to insert !!! Potentially the uploaded file is empty" - ) + log.debug("No chunks to insert !!! Potentially the uploaded file is empty") else: - log.info( - f"Vectordb insertion skipped (enable_insertion={self.enable_insertion})." - ) + log.info(f"Vectordb insertion skipped (enable_insertion={self.enable_insertion}).") # Mark task as completed await task_state_manager.set_state.remote(task_id, "COMPLETED") @@ -165,9 +155,7 @@ async def delete_file(self, file_id: str, partition: str) -> bool: try: await vectordb.delete_file.remote(file_id, partition) - log.info( - "Deleted file from partition.", file_id=file_id, partition=partition - ) + log.info("Deleted file from partition.", file_id=file_id, partition=partition) except Exception as e: log.exception("Error in delete_file", error=str(e)) @@ -177,16 +165,14 @@ async def delete_file(self, file_id: str, partition: str) -> bool: async def update_file_metadata( self, file_id: str, - metadata: Dict, + metadata: dict, partition: str, - user: Optional[Dict] = None, + user: dict | None = None, ): log = self.logger.bind(file_id=file_id, partition=partition) vectordb = ray.get_actor("Vectordb", namespace="openrag") if not self.enable_insertion: - log.error( - "Vector database is not enabled, but update_file_metadata was called." - ) + log.error("Vector database is not enabled, but update_file_metadata was called.") return try: @@ -206,16 +192,14 @@ async def update_file_metadata( async def copy_file( self, file_id: str, - metadata: Dict, + metadata: dict, partition: str, - user: Optional[Dict] = None, + user: dict | None = None, ): log = self.logger.bind(file_id=file_id, partition=partition) vectordb = ray.get_actor("Vectordb", namespace="openrag") if not self.enable_insertion: - log.error( - "Vector database is not enabled, but update_file_metadata was called." - ) + log.error("Vector database is not enabled, but update_file_metadata was called.") return try: @@ -242,10 +226,11 @@ async def asearch( query: str, top_k: int = 5, similarity_threshold: float = 0.80, - partition: Optional[Union[str, List[str]]] = None, - filter: Optional[Dict] = {}, - ) -> List[Document]: + partition: str | list[str] | None = None, + filter: dict | None = None, + ) -> list[Document]: partition_list = self._check_partition_list(partition) + filter = filter or {} vectordb = ray.get_actor("Vectordb", namespace="openrag") return await vectordb.async_search.remote( query=query, @@ -255,7 +240,7 @@ async def asearch( filter=filter, ) - def _check_partition_str(self, partition: Optional[str]) -> str: + def _check_partition_str(self, partition: str | None) -> str: if partition is None: self.logger.warning("partition not provided; using default.") return self.default_partition @@ -263,9 +248,7 @@ def _check_partition_str(self, partition: Optional[str]) -> str: raise ValueError("Partition must be a string.") return partition - def _check_partition_list( - self, partition: Optional[Union[str, List[str]]] - ) -> List[str]: + def _check_partition_list(self, partition: str | list[str] | None) -> list[str]: if partition is None: self.logger.warning("partition not provided; using default.") return [self.default_partition] @@ -278,17 +261,17 @@ def _check_partition_list( @dataclass class TaskInfo: - state: Optional[str] = None - error: Optional[str] = None - details: Dict[str, Any] = field(default_factory=dict) - object_ref: Optional[ray.ObjectRef] = None + state: str | None = None + error: str | None = None + details: dict[str, Any] = field(default_factory=dict) + object_ref: ray.ObjectRef | None = None @ray.remote(concurrency_groups={"set": 1000, "get": 1000, "queue_info": 1000}) class TaskStateManager: def __init__(self): - self.tasks: Dict[str, TaskInfo] = {} - self.user_index: Dict[int, set[str]] = {} + self.tasks: dict[str, TaskInfo] = {} + self.user_index: dict[int, set[str]] = {} self.lock = asyncio.Lock() async def _ensure_task(self, task_id: str) -> TaskInfo: @@ -336,36 +319,36 @@ async def set_object_ref(self, task_id: str, object_ref: dict): info.object_ref = object_ref @ray.method(concurrency_group="get") - async def get_state(self, task_id: str) -> Optional[str]: + async def get_state(self, task_id: str) -> str | None: async with self.lock: info = self.tasks.get(task_id) return info.state if info else None @ray.method(concurrency_group="get") - async def get_error(self, task_id: str) -> Optional[str]: + async def get_error(self, task_id: str) -> str | None: async with self.lock: info = self.tasks.get(task_id) return info.error if info else None @ray.method(concurrency_group="get") - async def get_details(self, task_id: str) -> Optional[dict]: + async def get_details(self, task_id: str) -> dict | None: async with self.lock: info = self.tasks.get(task_id) return info.details if info else None @ray.method(concurrency_group="get") - async def get_object_ref(self, task_id: str) -> Optional[dict]: + async def get_object_ref(self, task_id: str) -> dict | None: async with self.lock: info = self.tasks.get(task_id) return info.object_ref if info else None @ray.method(concurrency_group="queue_info") - async def get_all_states(self) -> Dict[str, str]: + async def get_all_states(self) -> dict[str, str]: async with self.lock: return {tid: info.state for tid, info in self.tasks.items()} @ray.method(concurrency_group="queue_info") - async def get_all_info(self) -> Dict[str, dict]: + async def get_all_info(self) -> dict[str, dict]: async with self.lock: return { task_id: { @@ -377,7 +360,7 @@ async def get_all_info(self) -> Dict[str, dict]: } @ray.method(concurrency_group="queue_info") - async def get_all_user_info(self, user_id: int) -> Dict[str, dict]: + async def get_all_user_info(self, user_id: int) -> dict[str, dict]: async with self.lock: task_ids = self.user_index.get(user_id, set()) return { @@ -391,7 +374,7 @@ async def get_all_user_info(self, user_id: int) -> Dict[str, dict]: } @ray.method(concurrency_group="queue_info") - async def get_pool_info(self) -> Dict[str, int]: + async def get_pool_info(self) -> dict[str, int]: return { "pool_size": POOL_SIZE, "max_tasks_per_worker": MAX_TASKS_PER_WORKER, diff --git a/openrag/components/indexer/loaders/__init__.py b/openrag/components/indexer/loaders/__init__.py index 2ec16bd60..7b142edce 100644 --- a/openrag/components/indexer/loaders/__init__.py +++ b/openrag/components/indexer/loaders/__init__.py @@ -7,7 +7,6 @@ import importlib import pkgutil from pathlib import Path -from typing import Dict, Set, Type from utils.logger import get_logger @@ -15,16 +14,15 @@ logger = get_logger() -def get_loader_classes(config: dict) -> Dict[str, Type[BaseLoader]]: + +def get_loader_classes(config: dict) -> dict[str, type[BaseLoader]]: # 1. Discover all subclasses root_pkg = "components.indexer.loaders" root_path = Path(__file__).parent - discovered: Dict[str, Type[BaseLoader]] = {} + discovered: dict[str, type[BaseLoader]] = {} - for finder, module_name, is_pkg in pkgutil.walk_packages( - path=[str(root_path)], prefix=f"{root_pkg}." - ): + for finder, module_name, is_pkg in pkgutil.walk_packages(path=[str(root_path)], prefix=f"{root_pkg}."): try: module = importlib.import_module(module_name) except ImportError as e: @@ -32,17 +30,13 @@ def get_loader_classes(config: dict) -> Dict[str, Type[BaseLoader]]: continue for attr in vars(module).values(): - if ( - isinstance(attr, type) - and issubclass(attr, BaseLoader) - and attr is not BaseLoader - ): + if isinstance(attr, type) and issubclass(attr, BaseLoader) and attr is not BaseLoader: discovered[attr.__name__] = attr # logger.debug(f"Discovered loaders: {discovered}") # 2. Read your config map of extensions → class names - loader_classes: Dict[str, Type[BaseLoader]] = {} + loader_classes: dict[str, type[BaseLoader]] = {} file_loaders = config.get("loader", {}).get("file_loaders", {}) for ext, cls_name in file_loaders.items(): @@ -57,7 +51,7 @@ def get_loader_classes(config: dict) -> Dict[str, Type[BaseLoader]]: return loader_classes -def get_supported_extensions(loader_classes: Dict[str, Type[BaseLoader]]) -> Set[str]: +def get_supported_extensions(loader_classes: dict[str, type[BaseLoader]]) -> set[str]: """ Get the set of supported file extensions from the loaded classes. diff --git a/openrag/components/indexer/loaders/base.py b/openrag/components/indexer/loaders/base.py index d095d262f..ee33ead41 100644 --- a/openrag/components/indexer/loaders/base.py +++ b/openrag/components/indexer/loaders/base.py @@ -3,16 +3,14 @@ from abc import ABC, abstractmethod from io import BytesIO from pathlib import Path -from typing import Dict, Optional, Union from components.prompts import IMAGE_DESCRIBER from components.utils import get_vlm_semaphore, load_config from langchain_core.messages import HumanMessage from langchain_openai import ChatOpenAI from PIL import Image -from utils.logger import get_logger - from utils.external_resource_errors import is_external_resource_error +from utils.logger import get_logger logger = get_logger() config = load_config() @@ -34,13 +32,12 @@ def __init__(self, **kwargs) -> None: self.image_captioning = self.config.loader.get("image_captioning", False) self.vlm_endpoint = ChatOpenAI(**settings).with_retry(stop_after_attempt=2) - self.min_width_pixels = 0 # minimum width in pixels - self.min_height_pixels = 0 # minimum height in pixels @abstractmethod async def aload_document( - file_path: Union[str, Path], - metadata: Optional[Dict] = None, + self, + file_path: str | Path, + metadata: dict | None = None, save_markdown: bool = False, ): pass @@ -68,7 +65,7 @@ def _is_data_uri(self, data: str) -> bool: async def get_image_description( self, - image_data: Union[Image.Image, str], + image_data: Image.Image | str, ) -> str: """ Creates a description for an image using the LLM model. @@ -88,17 +85,6 @@ async def get_image_description( # Determine the type of image data and create appropriate message content if isinstance(image_data, Image.Image): # logger.info("Processing PIL Image", img_size=str(image_data.size)) - # Handle PIL Image - width, height = image_data.size - - # Check minimum dimensions - if ( - width <= self.min_width_pixels - or height <= self.min_height_pixels - ): - logger.debug( - f"Image too small: {width}x{height}, skipping description" - ) # Convert PIL Image to base64 img_b64 = self._pil_image_to_base64(image_data) @@ -123,9 +109,7 @@ async def get_image_description( image_url = f"data:image/png;base64,{image_data}" logger.debug("Processing raw base64 string") except Exception: - logger.error( - f"Invalid image data type or format: {type(image_data)}" - ) + logger.error(f"Invalid image data type or format: {type(image_data)}") return """\n\nInvalid image data format\n\n""" else: logger.error(f"Unsupported image data type: {type(image_data)}") @@ -161,9 +145,7 @@ async def get_image_description( log_extra["url"] = str(image_data) logger.warning(log_msg, **log_extra) else: - logger.exception( - "Error while generating image description", error=str(e) - ) + logger.exception("Error while generating image description", error=str(e)) image_description = "" return f"""\n\n{image_description}\n\n""" diff --git a/openrag/components/indexer/loaders/eml_loader.py b/openrag/components/indexer/loaders/eml_loader.py index e5fe5981b..c23bcca72 100644 --- a/openrag/components/indexer/loaders/eml_loader.py +++ b/openrag/components/indexer/loaders/eml_loader.py @@ -1,24 +1,24 @@ -from pathlib import Path -import tempfile -import os -from PIL import Image +import datetime +import email import io +import os +import tempfile +from email.utils import parsedate_to_datetime +from pathlib import Path from langchain_core.documents.base import Document -from .base import BaseLoader +from PIL import Image + from . import get_loader_classes -from typing import Optional +from .base import BaseLoader -import datetime -import json -import email -from email.utils import parsedate_to_datetime def json_serial(obj): - if isinstance(obj, datetime.datetime): - serial = obj.isoformat() - return serial - + if isinstance(obj, datetime.datetime): + serial = obj.isoformat() + return serial + + class EmlLoader(BaseLoader): def __init__(self, **kwargs): super().__init__(**kwargs) @@ -27,208 +27,230 @@ def __init__(self, **kwargs): # Get available loaders for processing attachments self.loader_classes = get_loader_classes(config=self.config) - async def aload_document( - self, file_path, metadata: dict = None, save_markdown: bool = False - ): + async def aload_document(self, file_path, metadata: dict | None = None, save_markdown: bool = False): try: with open(file_path, "rb") as fhdl: raw_email = fhdl.read() # Parse email using standard email library email_msg = email.message_from_bytes(raw_email) - + # Extract email metadata email_data = { - 'header': { - 'subject': email_msg.get('subject', ''), - 'from': email_msg.get('from', ''), - 'to': email_msg.get('to', ''), - 'date': email_msg.get('date', ''), - 'message-id': email_msg.get('message-id', ''), + "header": { + "subject": email_msg.get("subject", ""), + "from": email_msg.get("from", ""), + "to": email_msg.get("to", ""), + "date": email_msg.get("date", ""), + "message-id": email_msg.get("message-id", ""), }, - 'body': [], - 'attachment': [] + "body": [], + "attachment": [], } - + # Parse date if available - if email_data['header']['date']: + if email_data["header"]["date"]: try: - email_data['header']['date'] = parsedate_to_datetime(email_data['header']['date']).isoformat() - except: + email_data["header"]["date"] = parsedate_to_datetime(email_data["header"]["date"]).isoformat() + except Exception: pass - + # Extract body content and attachments body_content = "" - + for part in email_msg.walk(): content_type = part.get_content_type() content_disposition = part.get_content_disposition() - - if content_disposition == 'attachment' or content_disposition == 'inline': + + if content_disposition == "attachment" or content_disposition == "inline": # Handle attachments filename = part.get_filename() if filename: payload = part.get_payload(decode=True) if payload: attachment_info = { - 'filename': filename, - 'content_type': content_type, - 'size': len(payload), - 'raw': payload + "filename": filename, + "content_type": content_type, + "size": len(payload), + "raw": payload, } - email_data['attachment'].append(attachment_info) - - elif content_type.startswith('text/'): + email_data["attachment"].append(attachment_info) + + elif content_type.startswith("text/"): # Handle text content - if content_type == 'text/plain' or content_type == 'text/html': + if content_type == "text/plain" or content_type == "text/html": text_content = part.get_payload(decode=True) if text_content: try: # Try to decode as UTF-8, fallback to latin-1 if isinstance(text_content, bytes): try: - text_content = text_content.decode('utf-8') + text_content = text_content.decode("utf-8") except UnicodeDecodeError: - text_content = text_content.decode('latin-1', errors='ignore') - + text_content = text_content.decode("latin-1", errors="ignore") + body_info = { - 'content': text_content, - 'content_type': content_type + "content": text_content, + "content_type": content_type, } - email_data['body'].append(body_info) - + email_data["body"].append(body_info) + # Use plain text as primary body content - if content_type == 'text/plain' or not body_content: + if content_type == "text/plain" or not body_content: body_content = text_content except Exception as e: print(f"Failed to decode text content: {e}") # Extract body content content_body = body_content.strip() if body_content else "" - + # Process attachments using appropriate loaders attachments_text = "" - if email_data['attachment']: + if email_data["attachment"]: attachments_text = "\n\n--- ATTACHMENTS ---\n" - for attachment in email_data['attachment']: - filename = attachment.get('filename', 'unknown') - content_type = attachment.get('content_type', 'unknown') - size = attachment.get('size', 'unknown') - + for attachment in email_data["attachment"]: + filename = attachment.get("filename", "unknown") + content_type = attachment.get("content_type", "unknown") + size = attachment.get("size", "unknown") + attachments_text += f"\nAttachment: {filename}\n" attachments_text += f"Content-Type: {content_type}\n" attachments_text += f"Size: {size} bytes\n" - + # Try to process attachment using appropriate loader - if 'raw' in attachment: + if "raw" in attachment: try: # Get file extension from filename file_ext = Path(filename).suffix.lower() - + # Check if we have a loader for this file type loader_cls = self.loader_classes.get(file_ext) - + if loader_cls: # Save attachment to temporary file with tempfile.NamedTemporaryFile(suffix=file_ext, delete=False) as temp_file: - temp_file.write(attachment['raw']) + temp_file.write(attachment["raw"]) temp_file_path = temp_file.name - + try: # Use appropriate loader to process attachment loader = loader_cls(**self.kwargs) attachment_doc = await loader.aload_document( - temp_file_path, - metadata={'source': f'attachment:{filename}'} + temp_file_path, + metadata={"source": f"attachment:{filename}"}, ) attachments_text += f"Content:\n{attachment_doc.page_content}\n" except Exception as e: attachments_text += f"Failed to process attachment with loader ({loader_cls.__name__}): {str(e)[:200]}...\n" - + # Special fallback handling for PDFs with alternative loaders - if file_ext == '.pdf': - pdf_fallback_loaders = ['PyMuPDFLoader', 'PyMuPDF4LLMLoader', 'DoclingLoader'] + if file_ext == ".pdf": + pdf_fallback_loaders = [ + "PyMuPDFLoader", + "PyMuPDF4LLMLoader", + "DoclingLoader", + ] fallback_success = False - + for fallback_loader_name in pdf_fallback_loaders: - if fallback_loader_name != loader_cls.__name__: # Don't try the same loader again + if ( + fallback_loader_name != loader_cls.__name__ + ): # Don't try the same loader again try: # Try to get the fallback loader class fallback_loader_cls = None - for ext, cls in self.loader_classes.items(): + for ( + ext, + cls, + ) in self.loader_classes.items(): if cls.__name__ == fallback_loader_name: fallback_loader_cls = cls break - + if fallback_loader_cls: - attachments_text += f"Trying fallback PDF loader: {fallback_loader_name}\n" + attachments_text += ( + f"Trying fallback PDF loader: {fallback_loader_name}\n" + ) fallback_loader = fallback_loader_cls(**self.kwargs) attachment_doc = await fallback_loader.aload_document( - temp_file_path, - metadata={'source': f'attachment:{filename}'} + temp_file_path, + metadata={"source": f"attachment:{filename}"}, ) attachments_text += f"Content (via {fallback_loader_name}):\n{attachment_doc.page_content}\n" fallback_success = True break except Exception as fallback_e: attachments_text += f"Fallback {fallback_loader_name} also failed: {str(fallback_e)[:100]}...\n" - + if not fallback_success: attachments_text += f"All PDF loaders failed for {filename}\n" - + # Try fallback processing for images - if file_ext in ['.png', '.jpg', '.jpeg', '.svg']: + if file_ext in [".png", ".jpg", ".jpeg", ".svg"]: try: if self.config.loader.get("image_captioning", False): # Try to load image directly from bytes as fallback - image = Image.open(io.BytesIO(attachment['raw'])) - caption = await self.get_image_description(image=image) + image = Image.open(io.BytesIO(attachment["raw"])) + caption = await self.get_image_description(image_data=image) attachments_text += f"Fallback Image Description:\n{caption}\n" else: - attachments_text += f"Image attachment present but image captioning disabled\n" + attachments_text += ( + "Image attachment present but image captioning disabled\n" + ) except Exception as img_e: attachments_text += f"Image fallback also failed: {str(img_e)[:100]}...\n" - + # Try text fallback for other text-based formats - elif file_ext in ['.txt', '.docx', '.doc'] or (file_ext == '.pdf' and not fallback_success): + elif file_ext in [".txt", ".docx", ".doc"] or ( + file_ext == ".pdf" and not fallback_success + ): try: # Try to extract any readable text directly - text_content = attachment['raw'].decode('utf-8', errors='ignore') + text_content = attachment["raw"].decode("utf-8", errors="ignore") if text_content.strip(): - attachments_text += f"Fallback text extraction:\n{text_content[:1000]}...\n" + attachments_text += ( + f"Fallback text extraction:\n{text_content[:1000]}...\n" + ) else: - attachments_text += f"No readable text found in attachment\n" + attachments_text += "No readable text found in attachment\n" except Exception as text_e: attachments_text += f"Text fallback failed: {str(text_e)[:100]}...\n" finally: # Clean up temporary file if os.path.exists(temp_file_path): os.unlink(temp_file_path) - + # Special handling for images with captioning if no specific loader or captioning is enabled - elif file_ext in ['.png', '.jpg', '.jpeg', '.svg'] and self.config.loader.get("image_captioning", False): + elif file_ext in [ + ".png", + ".jpg", + ".jpeg", + ".svg", + ] and self.config.loader.get("image_captioning", False): try: # Load image from raw bytes - image = Image.open(io.BytesIO(attachment['raw'])) + image = Image.open(io.BytesIO(attachment["raw"])) # Verify image can be processed image.verify() # Reopen image since verify() closes it - image = Image.open(io.BytesIO(attachment['raw'])) + image = Image.open(io.BytesIO(attachment["raw"])) # Generate caption using the base loader's method - caption = await self.get_image_description(image=image) + caption = await self.get_image_description(image_data=image) attachments_text += f"Image Description:\n{caption}\n" except Exception as e: attachments_text += f"Failed to generate image caption: {str(e)[:200]}...\n" # Try to show basic image info if available try: size_info = f"Image size: {len(attachment['raw'])} bytes" - attachments_text += f"Image attachment present but corrupted or unreadable. {size_info}\n" - except: - attachments_text += f"Image attachment present but corrupted or unreadable\n" - - elif content_type.startswith('text/'): + attachments_text += ( + f"Image attachment present but corrupted or unreadable. {size_info}\n" + ) + except Exception: + attachments_text += "Image attachment present but corrupted or unreadable\n" + + elif content_type.startswith("text/"): # For text attachments, decode directly - attachment_content = attachment['raw'].decode('utf-8', errors='ignore') + attachment_content = attachment["raw"].decode("utf-8", errors="ignore") attachments_text += f"Content:\n{attachment_content}\n" else: # For other binary content, just show metadata @@ -236,42 +258,46 @@ async def aload_document( except Exception as e: attachments_text += f"Content could not be processed: {e}\n" attachments_text += "---\n" - + # Combine body and attachments content_body = content_body + attachments_text # Prepare metadata if metadata is None: metadata = {} - + # Add email metadata to document metadata - metadata.update({ - 'email_subject': email_data['header']['subject'], - 'email_from': email_data['header']['from'], - 'email_to': email_data['header']['to'], - 'email_date': email_data['header']['date'], - 'email_message_id': email_data['header']['message-id'], - 'email_attachment_count': len(email_data['attachment']), - 'email_attachment_filenames': [att['filename'] for att in email_data['attachment']], - }) - + metadata.update( + { + "email_subject": email_data["header"]["subject"], + "email_from": email_data["header"]["from"], + "email_to": email_data["header"]["to"], + "email_date": email_data["header"]["date"], + "email_message_id": email_data["header"]["message-id"], + "email_attachment_count": len(email_data["attachment"]), + "email_attachment_filenames": [att["filename"] for att in email_data["attachment"]], + } + ) + # Add attachment metadata if there are attachments - if email_data['attachment']: + if email_data["attachment"]: attachment_metadata = [] - for att in email_data['attachment']: - attachment_metadata.append({ - 'filename': att['filename'], - 'content_type': att['content_type'], - 'size': att['size'] - }) - metadata['email_attachments'] = attachment_metadata + for att in email_data["attachment"]: + attachment_metadata.append( + { + "filename": att["filename"], + "content_type": att["content_type"], + "size": att["size"], + } + ) + metadata["email_attachments"] = attachment_metadata # Save content body to a file if requested if save_markdown: - markdown_path = Path(file_path).with_suffix('.md') - with open(markdown_path, 'w', encoding='utf-8') as md_file: + markdown_path = Path(file_path).with_suffix(".md") + with open(markdown_path, "w", encoding="utf-8") as md_file: md_file.write(content_body) - metadata['markdown_path'] = str(markdown_path) + metadata["markdown_path"] = str(markdown_path) except Exception as e: raise ValueError(f"Failed to parse the EML file {file_path}: {e}") diff --git a/openrag/components/indexer/loaders/image.py b/openrag/components/indexer/loaders/image.py index bd78f6605..9700f15bc 100644 --- a/openrag/components/indexer/loaders/image.py +++ b/openrag/components/indexer/loaders/image.py @@ -1,7 +1,9 @@ -from .base import BaseLoader -from PIL import Image from pathlib import Path + from langchain_core.documents import Document +from PIL import Image + +from .base import BaseLoader class ImageLoader(BaseLoader): diff --git a/openrag/components/indexer/loaders/pdf_loaders/docling2.py b/openrag/components/indexer/loaders/pdf_loaders/docling2.py index 6c6d9cc37..c964cab1a 100644 --- a/openrag/components/indexer/loaders/pdf_loaders/docling2.py +++ b/openrag/components/indexer/loaders/pdf_loaders/docling2.py @@ -1,5 +1,8 @@ import asyncio + +import ray import torch +from config import load_config from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend from docling.datamodel.base_models import InputFormat from docling.datamodel.document import ConversionResult @@ -15,13 +18,8 @@ from langchain_core.documents.base import Document from tqdm.asyncio import tqdm from utils.logger import get_logger -from config import load_config - -import asyncio from ..base import BaseLoader -import ray - logger = get_logger() config = load_config() diff --git a/openrag/components/indexer/loaders/pdf_loaders/marker.py b/openrag/components/indexer/loaders/pdf_loaders/marker.py index e7446d925..cfedd1520 100644 --- a/openrag/components/indexer/loaders/pdf_loaders/marker.py +++ b/openrag/components/indexer/loaders/pdf_loaders/marker.py @@ -3,7 +3,6 @@ import re import time from pathlib import Path -from typing import Dict, Optional, Union import ray import torch @@ -217,8 +216,8 @@ def __init__(self, **kwargs): async def aload_document( self, - file_path: Union[str, Path], - metadata: Optional[Dict] = None, + file_path: str | Path, + metadata: dict | None = None, save_markdown: bool = False, ) -> Document: from components.ray_utils import call_ray_actor_with_timeout diff --git a/openrag/components/indexer/loaders/pdf_loaders/openai.py b/openrag/components/indexer/loaders/pdf_loaders/openai.py index 62700b14b..06535ae57 100644 --- a/openrag/components/indexer/loaders/pdf_loaders/openai.py +++ b/openrag/components/indexer/loaders/pdf_loaders/openai.py @@ -5,7 +5,6 @@ import time from abc import ABC, abstractmethod from pathlib import Path -from typing import Dict, List, Optional, Union import pypdfium2 as pdfium from langchain.schema import Document @@ -16,7 +15,7 @@ from ..base import BaseLoader -async def pdf_to_images(self, pdf_path: str, scale: float = 1.0) -> List[Image.Image]: +async def pdf_to_images(pdf_path: str, scale: float = 1.0) -> list[Image.Image]: pdf: pdfium.PdfDocument = await asyncio.to_thread(pdfium.PdfDocument, pdf_path) return [p.render(scale=scale).to_pil() for p in pdf] @@ -38,14 +37,12 @@ def __init__(self, **kwargs): max_retries=self.config.loader["openai"].get("max_retries", 2), top_p=self.config.loader["openai"].get("top_p", 0.9), ) - self.llm_semaphore = asyncio.Semaphore( - self.config.loader["openai"].get("concurrency_limit", 20) - ) + self.llm_semaphore = asyncio.Semaphore(self.config.loader["openai"].get("concurrency_limit", 20)) async def aload_document( self, - file_path: Union[str, Path], - metadata: Optional[Dict] = None, + file_path: str | Path, + metadata: dict | None = None, save_markdown: bool = False, ) -> Document: """Main pipeline: PDF → OCR → Caption → Markdown.""" @@ -71,13 +68,11 @@ async def aload_document( logger.exception("Error in OpenAILoader.aload_document", path=file_path) raise - async def _run_ocr_on_pages(self, pages: List[Image.Image]) -> List[dict]: + async def _run_ocr_on_pages(self, pages: list[Image.Image]) -> list[dict]: tasks = [self._img2result(img) for img in pages] return await asyncio.gather(*tasks) - async def _assemble_markdown( - self, pages: List[Image.Image], results: List[dict] - ) -> str: + async def _assemble_markdown(self, pages: list[Image.Image], results: list[dict]) -> str: markdown_parts = [] for page_img, page_res in zip(pages, results): if not page_res: @@ -108,9 +103,7 @@ async def _img2result(self, img: Image.Image, format: str = "PNG") -> dict: "content": [ { "type": "image_url", - "image_url": { - "url": f"data:image/{format.lower()};base64,{img_b64}" - }, + "image_url": {"url": f"data:image/{format.lower()};base64,{img_b64}"}, }, { "type": "text", @@ -131,9 +124,7 @@ async def _img2result(self, img: Image.Image, format: str = "PNG") -> dict: @abstractmethod def _result_to_md(self, result: list[dict]) -> str: """Convert structured OCR + caption results to markdown format.""" - pass @abstractmethod async def _caption_images(self, page_img: Image.Image, page_res: list): """Extract picture elements and caption them.""" - pass diff --git a/openrag/components/indexer/loaders/serializer.py b/openrag/components/indexer/loaders/serializer.py index 9270c1b22..9f687a864 100644 --- a/openrag/components/indexer/loaders/serializer.py +++ b/openrag/components/indexer/loaders/serializer.py @@ -1,6 +1,5 @@ import gc from pathlib import Path -from typing import Dict, Optional, Union import ray import torch @@ -40,9 +39,10 @@ def __init__(self, data_dir=None, **kwargs) -> None: async def serialize_document( self, task_id: str, - path: Union[str, Path], - metadata: Optional[Dict] = {}, + path: str | Path, + metadata: dict | None = None, ) -> Document: + metadata = metadata or {} # Set task state log = self.logger.bind( file_id=metadata.get("file_id"), diff --git a/openrag/components/indexer/loaders/test_media_loader.py b/openrag/components/indexer/loaders/test_media_loader.py index 07273e5ec..951fa618a 100644 --- a/openrag/components/indexer/loaders/test_media_loader.py +++ b/openrag/components/indexer/loaders/test_media_loader.py @@ -249,7 +249,9 @@ def test_short_audio_single_chunk(self): def test_long_audio_multiple_chunks(self): """Test that long audio is split into multiple chunks.""" # Create 10 seconds of audio with silences - segment = Sine(440).to_audio_segment(duration=2000) + AudioSegment.silent(duration=500) + segment = Sine(440).to_audio_segment(duration=2000) + AudioSegment.silent( + duration=500 + ) audio = segment * 3 # Repeat pattern chunks = self.get_audio_chunks( diff --git a/openrag/components/indexer/loaders/txt_loader.py b/openrag/components/indexer/loaders/txt_loader.py index aa9b79100..3c9efad5c 100644 --- a/openrag/components/indexer/loaders/txt_loader.py +++ b/openrag/components/indexer/loaders/txt_loader.py @@ -2,18 +2,15 @@ Text and Markdown file loader implementation. """ -import asyncio +import itertools import re from pathlib import Path -from typing import Dict, Optional, Union from components.indexer.loaders.base import BaseLoader from langchain_community.document_loaders import TextLoader as LangchainTextLoader from langchain_core.documents.base import Document -from utils.logger import get_logger -import itertools from tqdm.asyncio import tqdm - +from utils.logger import get_logger logger = get_logger() @@ -28,8 +25,8 @@ def __init__(self, **kwargs) -> None: async def aload_document( self, - file_path: Union[str, Path], - metadata: Optional[Dict] = None, + file_path: str | Path, + metadata: dict | None = None, save_markdown: bool = False, ) -> Document: if metadata is None: @@ -67,8 +64,8 @@ def __init__(self, **kwargs) -> None: async def aload_document( self, - file_path: Union[str, Path], - metadata: Optional[Dict] = None, + file_path: str | Path, + metadata: dict | None = None, save_markdown: bool = False, ) -> Document: if metadata is None: diff --git a/openrag/components/indexer/utils/files.py b/openrag/components/indexer/utils/files.py index da685f29d..f9194de46 100644 --- a/openrag/components/indexer/utils/files.py +++ b/openrag/components/indexer/utils/files.py @@ -2,7 +2,6 @@ import secrets import time from pathlib import Path -from typing import Dict, Optional import aiofiles import consts @@ -71,11 +70,12 @@ async def save_file_to_disk( return file_path -async def serialize_file(task_id: str, path: str, metadata: Optional[Dict] = {}): +async def serialize_file(task_id: str, path: str, metadata: dict | None = None): import ray - from components.ray_utils import call_ray_actor_with_timeout + metadata = metadata or {} + serializer = ray.get_actor("DocSerializer", namespace="openrag") future = serializer.serialize_document.remote(task_id, path, metadata=metadata) diff --git a/openrag/components/indexer/utils/test_text_sanitizer.py b/openrag/components/indexer/utils/test_text_sanitizer.py index a6711e5d9..d26f2be22 100644 --- a/openrag/components/indexer/utils/test_text_sanitizer.py +++ b/openrag/components/indexer/utils/test_text_sanitizer.py @@ -124,10 +124,7 @@ def test_extra_spaces_in_cells(self): table = "| Header 1 | Header 2 |\n|-------------|-------------|\n| Cell 1 | Cell 2 |" result = clean_markdown_table_spacing(table) - assert ( - result - == "| Header 1 | Header 2 |\n| ------------- | ------------- |\n| Cell 1 | Cell 2 |" - ) + assert result == "| Header 1 | Header 2 |\n| ------------- | ------------- |\n| Cell 1 | Cell 2 |" def test_inconsistent_spacing(self): """Test normalizing inconsistent spacing across rows.""" @@ -141,19 +138,14 @@ def test_empty_cells(self): table = "| Col1 | Col2 | Col3 |\n|------|------|------|\n| Data | | More |\n| | Data | |" result = clean_markdown_table_spacing(table) - assert ( - result - == "| Col1 | Col2 | Col3 |\n| ------ | ------ | ------ |\n| Data | | More |\n| | Data | |" - ) + assert result == "| Col1 | Col2 | Col3 |\n| ------ | ------ | ------ |\n| Data | | More |\n| | Data | |" def test_multiline_spacing(self): """Test table with varying amounts of whitespace.""" table = "| A | B | C |\n|------|--------|----------|\n|1|2|3|" result = clean_markdown_table_spacing(table) - assert ( - result == "| A | B | C |\n| ------ | -------- | ---------- |\n| 1 | 2 | 3 |" - ) + assert result == "| A | B | C |\n| ------ | -------- | ---------- |\n| 1 | 2 | 3 |" class TestSanitizeExtractedText: @@ -192,7 +184,7 @@ def test_basic_extraction_scenario(self): Some \x00control\x01 characters\x02 too. And zero-width\u200bspaces. - """ + """ # noqa: W293 result = sanitize_extracted_text(text) diff --git a/openrag/components/indexer/vectordb/utils.py b/openrag/components/indexer/vectordb/utils.py index 99c63c40b..de6c67f75 100644 --- a/openrag/components/indexer/vectordb/utils.py +++ b/openrag/components/indexer/vectordb/utils.py @@ -2,7 +2,6 @@ import os import secrets from datetime import datetime -from typing import Dict, Optional from sqlalchemy import ( JSON, @@ -153,7 +152,7 @@ def __init__(self, database_url: str, logger=logger): except Exception as e: raise VDBConnectionError( - f"Failed to connect to database: {str(e)}", + f"Failed to connect to database: {e!s}", db_url=database_url, db_type="SQLAlchemy", ) @@ -179,7 +178,7 @@ def _ensure_admin_user(self, admin_token: str): s.commit() self.logger.info("Upgraded existing user to admin") - def list_partition_files(self, partition: str, limit: Optional[int] = None): + def list_partition_files(self, partition: str, limit: int | None = None): """List files in a partition with optional limit - Optimized by querying File table directly""" log = self.logger.bind(partition=partition) with self.Session() as session: @@ -208,8 +207,8 @@ def add_file_to_partition( self, file_id: str, partition: str, - file_metadata: Optional[Dict] = None, - user_id: Optional[int] = None, + file_metadata: dict | None = None, + user_id: int | None = None, ): """Add a file to a partition - Optimized with direct partition lookup""" log = self.logger.bind(file_id=file_id, partition=partition) @@ -334,8 +333,8 @@ def file_exists_in_partition(self, file_id: str, partition: str): def create_user( self, - display_name: Optional[str] = None, - external_user_id: Optional[str] = None, + display_name: str | None = None, + external_user_id: str | None = None, is_admin: bool = False, ) -> dict: """Create a user and generate an API token for them.""" @@ -375,7 +374,7 @@ def list_users(self) -> list[dict]: for u in users ] - def get_user_by_token(self, token: str) -> Optional[dict]: + def get_user_by_token(self, token: str) -> dict | None: with self.Session() as s: hashed_token = self.hash_token(token) user = s.query(User).filter(User.token == hashed_token).first() @@ -399,7 +398,7 @@ def get_user_by_token(self, token: str) -> Optional[dict]: "memberships": memberships, } - def get_user_by_id(self, user_id: int) -> Optional[dict]: + def get_user_by_id(self, user_id: int) -> dict | None: with self.Session() as s: user = s.query(User).filter(User.id == user_id).first() if not user: diff --git a/openrag/components/indexer/vectordb/vectordb.py b/openrag/components/indexer/vectordb/vectordb.py index f6fd6c367..291f8ac04 100644 --- a/openrag/components/indexer/vectordb/vectordb.py +++ b/openrag/components/indexer/vectordb/vectordb.py @@ -1,7 +1,6 @@ import asyncio import time from abc import ABC, abstractmethod -from typing import List, Optional import numpy as np import ray @@ -55,7 +54,7 @@ async def delete_partition(self, partition: str): pass @abstractmethod - def list_partition_files(self, partition: str, limit: Optional[int] = None): + def list_partition_files(self, partition: str, limit: int | None = None): pass @abstractmethod @@ -73,7 +72,7 @@ async def async_search( top_k: int = 5, similarity_threshold: int = 0.60, partition: list[str] = None, - filter: Optional[dict] = None, + filter: dict | None = None, with_surrounding_chunks: bool = False, ) -> list[Document]: pass @@ -85,21 +84,17 @@ async def async_multi_query_search( queries: list[str], top_k_per_query: int = 5, similarity_threshold: int = 0.6, - filter: Optional[dict] = None, + filter: dict | None = None, with_surrounding_chunks: bool = False, ) -> list[Document]: pass @abstractmethod - async def list_all_chunk( - self, partition: str, include_embedding: bool = True - ) -> List[Document]: + async def list_all_chunk(self, partition: str, include_embedding: bool = True) -> list[Document]: pass @abstractmethod - async def get_file_chunks( - self, file_id: str, partition: str, include_id: bool = False, limit: int = 100 - ): + async def get_file_chunks(self, file_id: str, partition: str, include_id: bool = False, limit: int = 100): pass @abstractmethod @@ -155,15 +150,13 @@ def __init__(self): self._async_client = AsyncMilvusClient(uri=uri) except MilvusException as e: raise VDBConnectionError( - f"Failed to connect to Milvus: {str(e)}", + f"Failed to connect to Milvus: {e!s}", db_url=uri, db_type="Milvus", ) # embedder - self.embedder: BaseEmbedding = EmbeddingFactory.get_embedder( - embeddings_config=self.config.embedder - ) + self.embedder: BaseEmbedding = EmbeddingFactory.get_embedder(embeddings_config=self.config.embedder) self.hybrid_search = self.config.vectordb.get("hybrid_search", True) # partition related params @@ -174,9 +167,7 @@ def __init__(self): self.partition_file_manager: PartitionFileManager = None # Initialize collection-related attributes - self.collection_name = self.config.vectordb.get( - "collection_name", "vdb_test" - ) + self.collection_name = self.config.vectordb.get("collection_name", "vdb_test") self.collection_loaded = False self.load_collection() @@ -184,25 +175,19 @@ def __init__(self): raise except Exception as e: - self.logger.exception( - "Unexpected error initializing Milvus clients", error=str(e) - ) + self.logger.exception("Unexpected error initializing Milvus clients", error=str(e)) raise VDBConnectionError( - f"Unexpected error initializing Milvus clients: {str(e)}", + f"Unexpected error initializing Milvus clients: {e!s}", db_url=uri, db_type="Milvus", ) def load_collection(self): if not self.collection_loaded: - self.logger = self.logger.bind( - collection=self.collection_name, database="Milvus" - ) + self.logger = self.logger.bind(collection=self.collection_name, database="Milvus") try: if self._client.has_collection(self.collection_name): - self.logger.warning( - f"Collection `{self.collection_name}` already exists. Loading it." - ) + self.logger.warning(f"Collection `{self.collection_name}` already exists. Loading it.") else: self.logger.info("Creating empty collection") index_params = self._create_index() @@ -222,7 +207,7 @@ def load_collection(self): error=str(e), ) raise VDBCreateOrLoadCollectionError( - f"Failed to create collection `{self.collection_name}`: {str(e)}", + f"Failed to create collection `{self.collection_name}`: {e!s}", collection_name=self.collection_name, operation="create_collection", ) @@ -235,7 +220,7 @@ def load_collection(self): error=str(e), ) raise VDBCreateOrLoadCollectionError( - f"Failed to load existing collection `{self.collection_name}`: {str(e)}", + f"Failed to load existing collection `{self.collection_name}`: {e!s}", collection_name=self.collection_name, operation="load_collection", ) @@ -253,16 +238,14 @@ def load_collection(self): error=str(e), ) raise UnexpectedVDBError( - f"Unexpected error setting collection name `{self.collection_name}`: {str(e)}", + f"Unexpected error setting collection name `{self.collection_name}`: {e!s}", collection_name=self.collection_name, ) def _create_schema(self): self.logger.info("Creating Schema") schema = self._client.create_schema(enable_dynamic_field=True) - schema.add_field( - field_name="_id", datatype=DataType.INT64, is_primary=True, auto_id=True - ) + schema.add_field(field_name="_id", datatype=DataType.INT64, is_primary=True, auto_id=True) schema.add_field( field_name="text", datatype=DataType.VARCHAR, @@ -322,9 +305,7 @@ def _create_index(self): ) # ADD index for partition field - index_params.add_index( - field_name="partition", index_type="INVERTED", index_name="partition_idx" - ) + index_params.add_index(field_name="partition", index_type="INVERTED", index_name="partition_idx") # Add index for vector field index_params.add_index( @@ -368,9 +349,7 @@ async def async_add_documents(self, chunks: list[Document], user: dict) -> None: ) # check if this file_id exists - res = self.partition_file_manager.file_exists_in_partition( - file_id=file_id, partition=partition - ) + res = self.partition_file_manager.file_exists_in_partition(file_id=file_id, partition=partition) if res: error_msg = f"This File Id ({file_id}) already exists in Partition ({partition})" self.logger.error(error_msg) @@ -416,11 +395,9 @@ async def async_add_documents(self, chunks: list[Document], user: dict) -> None: raise except Exception as e: - self.logger.exception( - "Unexpected error while adding a document", error=str(e) - ) + self.logger.exception("Unexpected error while adding a document", error=str(e)) raise UnexpectedVDBError( - f"Unexpected error while adding a document: {str(e)}", + f"Unexpected error while adding a document: {e!s}", collection_name=self.collection_name, ) @@ -460,7 +437,7 @@ async def async_search( top_k: int = 5, similarity_threshold: int = 0.80, partition: list[str] = None, - filter: Optional[dict] = None, + filter: dict | None = None, with_surrounding_chunks: bool = False, ) -> list[Document]: expr_parts = [] @@ -524,9 +501,7 @@ async def async_search( if with_surrounding_chunks: self.logger.debug("Fetching surrounding chunks") surrounding_chunks = await self.get_surrounding_chunks(docs) - self.logger.debug( - "Fetched surrounding chunks", count=len(surrounding_chunks) - ) + self.logger.debug("Fetched surrounding chunks", count=len(surrounding_chunks)) docs.extend(surrounding_chunks) return docs @@ -534,26 +509,24 @@ async def async_search( except MilvusException as e: self.logger.exception("Search failed in Milvus", error=str(e)) raise VDBSearchError( - f"Search failed in Milvus: {str(e)}", + f"Search failed in Milvus: {e!s}", collection_name=self.collection_name, partition=partition, ) except EmbeddingError as e: - self.logger.exception( - "Embedding failed while processing the query", error=str(e) - ) + self.logger.exception("Embedding failed while processing the query", error=str(e)) raise except Exception as e: self.logger.exception("Unexpected error occurred", error=str(e)) raise UnexpectedVDBError( - f"Unexpected error occurred: {str(e)}", + f"Unexpected error occurred: {e!s}", collection_name=self.collection_name, partition=partition, ) async def get_surrounding_chunks(self, docs: list[Document]) -> list[Document]: - existant_ids = set(doc.metadata.get("_id") for doc in docs) + existant_ids = {doc.metadata.get("_id") for doc in docs} # Collect all prev/next section IDs section_ids = [ @@ -591,11 +564,7 @@ async def get_surrounding_chunks(self, docs: list[Document]) -> list[Document]: output_docs.append( Document( page_content=response[0]["text"], - metadata={ - key: value - for key, value in response[0].items() - if key not in ["text", "vector"] - }, + metadata={key: value for key, value in response[0].items() if key not in ["text", "vector"]}, ) ) @@ -609,19 +578,13 @@ async def delete_file(self, file_id: str, partition: str): filter=f"partition == '{partition}' and file_id == '{file_id}'", ) - self.partition_file_manager.remove_file_from_partition( - file_id=file_id, partition=partition - ) - log.info( - "Deleted file chunks from partition.", count=res.get("delete_count", 0) - ) + self.partition_file_manager.remove_file_from_partition(file_id=file_id, partition=partition) + log.info("Deleted file chunks from partition.", count=res.get("delete_count", 0)) except MilvusException as e: - log.exception( - f"Couldn't delete file chunks for file_id {file_id}", error=str(e) - ) + log.exception(f"Couldn't delete file chunks for file_id {file_id}", error=str(e)) raise VDBDeleteError( - f"Couldn't delete file chunks for file_id {file_id}: {str(e)}", + f"Couldn't delete file chunks for file_id {file_id}: {e!s}", collection_name=self.collection_name, partition=partition, file_id=file_id, @@ -631,15 +594,13 @@ async def delete_file(self, file_id: str, partition: str): except Exception as e: log.exception("Unexpected error while deleting file chunks", error=str(e)) raise UnexpectedVDBError( - f"Unexpected error while deleting file chunks {file_id}: {str(e)}", + f"Unexpected error while deleting file chunks {file_id}: {e!s}", collection_name=self.collection_name, partition=partition, file_id=file_id, ) - async def get_file_chunks( - self, file_id: str, partition: str, include_id: bool = False, limit: int = 100 - ): + async def get_file_chunks(self, file_id: str, partition: str, include_id: bool = False, limit: int = 100): log = self.logger.bind(file_id=file_id, partition=partition) try: self._check_file_exists(file_id, partition) @@ -650,9 +611,7 @@ async def get_file_chunks( # Pagination parameters offset = 0 results = [] - excluded_keys = ( - ["text", "vector", "_id"] if not include_id else ["text", "vector"] - ) + excluded_keys = ["text", "vector", "_id"] if not include_id else ["text", "vector"] while True: response = await self._async_client.query( @@ -672,11 +631,7 @@ async def get_file_chunks( docs = [ Document( page_content=res["text"], - metadata={ - key: value - for key, value in res.items() - if key not in excluded_keys - }, + metadata={key: value for key, value in res.items() if key not in excluded_keys}, ) for res in results ] @@ -684,11 +639,9 @@ async def get_file_chunks( return docs except MilvusException as e: - log.exception( - f"Couldn't get file chunks for file_id {file_id}", error=str(e) - ) + log.exception(f"Couldn't get file chunks for file_id {file_id}", error=str(e)) raise VDBSearchError( - f"Couldn't get file chunks for file_id {file_id}: {str(e)}", + f"Couldn't get file chunks for file_id {file_id}: {e!s}", collection_name=self.collection_name, partition=partition, file_id=file_id, @@ -699,7 +652,7 @@ async def get_file_chunks( except Exception as e: log.exception("Unexpected error while getting file chunks", error=str(e)) raise VDBSearchError( - f"Unexpected error while getting file chunks {file_id}: {str(e)}", + f"Unexpected error while getting file chunks {file_id}: {e!s}", collection_name=self.collection_name, partition=partition, file_id=file_id, @@ -730,24 +683,20 @@ async def get_chunk_by_id(self, chunk_id: str): if response: return Document( page_content=response[0]["text"], - metadata={ - key: value - for key, value in response[0].items() - if key not in ["text", "vector"] - }, + metadata={key: value for key, value in response[0].items() if key not in ["text", "vector"]}, ) return None except MilvusException as e: log.exception("Milvus query failed", error=str(e)) raise VDBSearchError( - f"Milvus query failed: {str(e)}", + f"Milvus query failed: {e!s}", collection_name=self.collection_name, ) except Exception as e: log.exception("Unexpected error while retrieving chunk", error=str(e)) raise UnexpectedVDBError( - f"Unexpected error while retrieving chunk {chunk_id}: {str(e)}", + f"Unexpected error while retrieving chunk {chunk_id}: {e!s}", collection_name=self.collection_name, ) @@ -756,9 +705,7 @@ def file_exists(self, file_id: str, partition: str): Check if a file exists in Milvus """ try: - return self.partition_file_manager.file_exists_in_partition( - file_id=file_id, partition=partition - ) + return self.partition_file_manager.file_exists_in_partition(file_id=file_id, partition=partition) except Exception as e: self.logger.exception( "File existence check failed.", @@ -768,12 +715,10 @@ def file_exists(self, file_id: str, partition: str): ) return False - def list_partition_files(self, partition: str, limit: Optional[int] = None): + def list_partition_files(self, partition: str, limit: int | None = None): try: self._check_partition_exists(partition) - return self.partition_file_manager.list_partition_files( - partition=partition, limit=limit - ) + return self.partition_file_manager.list_partition_files(partition=partition, limit=limit) except VDBError: raise @@ -784,7 +729,7 @@ def list_partition_files(self, partition: str, limit: Optional[int] = None): error=str(e), ) raise UnexpectedVDBError( - f"Unexpected error while listing files in partition {partition}: {str(e)}", + f"Unexpected error while listing files in partition {partition}: {e!s}", collection_name=self.collection_name, partition=partition, ) @@ -818,7 +763,7 @@ async def delete_partition(self, partition: str): except MilvusException as e: log.exception("Failed to delete partition", error=str(e)) raise VDBDeleteError( - f"Failed to delete partition `{partition}`: {str(e)}", + f"Failed to delete partition `{partition}`: {e!s}", collection_name=self.collection_name, partition=partition, ) @@ -828,7 +773,7 @@ async def delete_partition(self, partition: str): except Exception as e: log.exception("Unexpected error while deleting partition", error=str(e)) raise UnexpectedVDBError( - f"Unexpected error while deleting partition {partition}: {str(e)}", + f"Unexpected error while deleting partition {partition}: {e!s}", collection_name=self.collection_name, partition=partition, ) @@ -897,7 +842,7 @@ def prepare_metadata(res: dict): except MilvusException as e: self.logger.exception("Milvus query failed", error=str(e)) raise VDBSearchError( - f"Milvus query failed: {str(e)}", + f"Milvus query failed: {e!s}", collection_name=self.collection_name, partition=partition, ) @@ -910,7 +855,7 @@ def prepare_metadata(res: dict): error=str(e), ) raise UnexpectedVDBError( - f"Unexpected error while listing all chunks in partition {partition}: {str(e)}", + f"Unexpected error while listing all chunks in partition {partition}: {e!s}", collection_name=self.collection_name, partition=partition, ) @@ -921,9 +866,7 @@ async def create_user( external_user_id: str | None = None, is_admin: bool = False, ): - return self.partition_file_manager.create_user( - display_name, external_user_id, is_admin - ) + return self.partition_file_manager.create_user(display_name, external_user_id, is_admin) async def get_user(self, user_id: int): self._check_user_exists(user_id) @@ -932,9 +875,7 @@ async def get_user(self, user_id: int): async def delete_user(self, user_id: int): self._check_user_exists(user_id) user_partitions = [ - p["partition"] - for p in self.partition_file_manager.list_user_partitions(user_id) - if p["role"] == "owner" + p["partition"] for p in self.partition_file_manager.list_user_partitions(user_id) if p["role"] == "owner" ] for partition in user_partitions: self.partition_file_manager.delete_partition(partition) @@ -954,20 +895,14 @@ async def list_user_partitions(self, user_id: int): self._check_user_exists(user_id) return self.partition_file_manager.list_user_partitions(user_id) - async def list_partition_members(self, partition: str) -> List[dict]: + async def list_partition_members(self, partition: str) -> list[dict]: self._check_partition_exists(partition) return self.partition_file_manager.list_partition_members(partition) - async def update_partition_member_role( - self, partition: str, user_id: int, new_role: str - ): + async def update_partition_member_role(self, partition: str, user_id: int, new_role: str): self._check_membership_exists(partition, user_id) - self.partition_file_manager.update_partition_member_role( - partition, user_id, new_role - ) - self.logger.info( - f"User_id {user_id} role updated to '{new_role}' in partition '{partition}'." - ) + self.partition_file_manager.update_partition_member_role(partition, user_id, new_role) + self.logger.info(f"User_id {user_id} role updated to '{new_role}' in partition '{partition}'.") async def create_partition(self, partition: str, user_id: int): self._check_user_exists(user_id) @@ -1015,9 +950,7 @@ def _check_membership_exists(self, partition: str, user_id: int): ) def _check_file_exists(self, file_id, partition: str): - if not self.partition_file_manager.file_exists_in_partition( - file_id=file_id, partition=partition - ): + if not self.partition_file_manager.file_exists_in_partition(file_id=file_id, partition=partition): raise VDBFileNotFoundError( f"File ID '{file_id}' does not exist in partition '{partition}'", collection_name=self.collection_name, diff --git a/openrag/components/llm.py b/openrag/components/llm.py index 1eefd2955..634663c4a 100644 --- a/openrag/components/llm.py +++ b/openrag/components/llm.py @@ -38,9 +38,7 @@ async def completions(self, request: dict): yield data except httpx.HTTPStatusError as e: error_detail = e.response.text - raise ValueError( - f"LLM API error ({e.response.status_code}): {error_detail}" - ) + raise ValueError(f"LLM API error ({e.response.status_code}): {error_detail}") except json.JSONDecodeError as e: raise ValueError(f"Invalid JSON in API response: {str(e)}") @@ -63,9 +61,7 @@ async def chat_completion(self, request: dict): if response.status_code >= 400: await response.aread() error_detail = response.text - raise ValueError( - f"LLM API error ({response.status_code}): {error_detail}" - ) + raise ValueError(f"LLM API error ({response.status_code}): {error_detail}") async for line in response.aiter_lines(): yield line except ValueError: @@ -86,8 +82,6 @@ async def chat_completion(self, request: dict): yield data except httpx.HTTPStatusError as e: error_detail = e.response.text - raise ValueError( - f"LLM API error ({e.response.status_code}): {error_detail}" - ) + raise ValueError(f"LLM API error ({e.response.status_code}): {error_detail}") except json.JSONDecodeError as e: raise ValueError(f"Invalid JSON in API response: {str(e)}") diff --git a/openrag/components/map_reduce.py b/openrag/components/map_reduce.py index fb7eca0c3..e9f89f348 100644 --- a/openrag/components/map_reduce.py +++ b/openrag/components/map_reduce.py @@ -30,9 +30,7 @@ class SummarizedChunk(BaseModel): - relevancy: bool = Field( - ..., description="Indicates if the chunk is relevant to the query" - ) + relevancy: bool = Field(..., description="Indicates if the chunk is relevant to the query") summary: str = Field( "", description="The summarized content of the chunk. The field should be empty if relevancy is False.", @@ -51,9 +49,7 @@ class SummarizedChunk(BaseModel): class RAGMapReduce: def __init__(self, config): self.config = config - self.slm: ChatOpenAI = ChatOpenAI(**config.llm).with_structured_output( - SummarizedChunk - ) + self.slm: ChatOpenAI = ChatOpenAI(**config.llm).with_structured_output(SummarizedChunk) map_reduce_config = self.config.map_reduce self.initial_batch_size = map_reduce_config["initial_batch_size"] self.expansion_batch_size = map_reduce_config["expansion_batch_size"] @@ -77,9 +73,7 @@ async def infer_chunk_relevancy(self, query, chunk: Document) -> SummarizedChunk {"role": "system", "content": system_prompt_map}, { "role": "user", - "content": user_prompt.format( - query=query, content=chunk.page_content - ), + "content": user_prompt.format(query=query, content=chunk.page_content), }, ], **params, @@ -97,29 +91,23 @@ async def map_batch( kth_batch=1, ): """Process a batch of chunks""" - logger.debug( - f"Processing {kth_batch}-th batch of chunks", batch_size=len(chunks) - ) + logger.debug(f"Processing {kth_batch}-th batch of chunks", batch_size=len(chunks)) tasks = [self.infer_chunk_relevancy(query, chunk) for chunk in chunks] outputs: list[SummarizedChunk] = await tqdm.gather( *tasks, desc="Map & Reduce processing chunks", total=len(chunks) ) - terminate = all( - [not o.relevancy for o in outputs[-self.expansion_batch_size :]] - ) # if the last 'expansion_batch_size' chunks are all irrelevant, we can terminate + + # if the last 'expansion_batch_size' chunks are all irrelevant, we can terminate + terminate = all(not o.relevancy for o in outputs[-self.expansion_batch_size :]) for o, chunk in zip(outputs, chunks): if o.relevancy: - summaries.append( - Document(page_content=o.summary, metadata=chunk.metadata) - ) + summaries.append(Document(page_content=o.summary, metadata=chunk.metadata)) if self.debug: with open(LOG_DIR / "map_reduce.md", "a") as f: f.write(f"### Query: \n{query}\n") - f.write( - f"### Chunk Content: \n* Relevancy: {o.relevancy} \n\n {chunk.page_content}\n" - ) + f.write(f"### Chunk Content: \n* Relevancy: {o.relevancy} \n\n {chunk.page_content}\n") f.write(f"### Summary: \n{o.summary}\n") f.write("\n-------\n\n") @@ -142,23 +130,13 @@ async def map(self, query: str, chunks: list[Document]): chunks[: self.initial_batch_size], chunks[self.initial_batch_size :], ) - _, terminate = await self.map_batch( - query, initial_batch, summaries=summaries, kth_batch=1 - ) + _, terminate = await self.map_batch(query, initial_batch, summaries=summaries, kth_batch=1) - if ( - terminate - or not remaining_chunks - or len(summaries) >= self.max_total_documents - ): + if terminate or not remaining_chunks or len(summaries) >= self.max_total_documents: return summaries - for jth_batch, i in enumerate( - range(0, len(remaining_chunks), self.expansion_batch_size), start=2 - ): - n = min( - self.expansion_batch_size, self.max_total_documents - len(summaries) - ) + for jth_batch, i in enumerate(range(0, len(remaining_chunks), self.expansion_batch_size), start=2): + n = min(self.expansion_batch_size, self.max_total_documents - len(summaries)) if n <= 0: break @@ -175,7 +153,5 @@ async def map(self, query: str, chunks: list[Document]): if terminate or len(summaries) >= self.max_total_documents: break - logger.debug( - "Map reduce completed", relevant_chunks_count=len(summaries), query=query - ) + logger.debug("Map reduce completed", relevant_chunks_count=len(summaries), query=query) return summaries diff --git a/openrag/components/pipeline.py b/openrag/components/pipeline.py index bbfc3c1fa..a7fa3948f 100644 --- a/openrag/components/pipeline.py +++ b/openrag/components/pipeline.py @@ -204,7 +204,7 @@ async def completions(self, partition: list[str], payload: dict): llm_output = self.llm_client.completions(request=payload) return llm_output, docs except Exception as e: - logger.error(f"Error during chat completion: {str(e)}") + logger.error(f"Error during chat completion: {e!s}") raise e async def chat_completion(self, partition: list[str] | None, payload: dict): @@ -218,5 +218,5 @@ async def chat_completion(self, partition: list[str] | None, payload: dict): llm_output = self.llm_client.chat_completion(request=payload) return llm_output, docs except Exception as e: - logger.error(f"Error during chat completion: {str(e)}") + logger.error(f"Error during chat completion: {e!s}") raise e diff --git a/openrag/components/prompts/prompts.py b/openrag/components/prompts/prompts.py index 0f435368a..e7cf0ec6d 100644 --- a/openrag/components/prompts/prompts.py +++ b/openrag/components/prompts/prompts.py @@ -22,7 +22,7 @@ def load_prompt( if not file_path.exists(): raise FileNotFoundError(f"Prompt file not found: `{file_path}`") - with open(file_path, mode="r") as f: + with open(file_path) as f: sys_msg = f.read() return sys_msg diff --git a/openrag/components/reranker.py b/openrag/components/reranker.py index 474dfd576..178502751 100644 --- a/openrag/components/reranker.py +++ b/openrag/components/reranker.py @@ -1,9 +1,11 @@ import asyncio + from infinity_client import Client from infinity_client.api.default import rerank from infinity_client.models import RerankInput, ReRankResult from langchain_core.documents.base import Document + class Reranker: def __init__(self, logger, config): self.model_name = config.reranker["model_name"] diff --git a/openrag/components/retriever.py b/openrag/components/retriever.py index 43d4deea5..3c5caf32a 100644 --- a/openrag/components/retriever.py +++ b/openrag/components/retriever.py @@ -1,5 +1,6 @@ # Import necessary modules and classes from abc import ABC, abstractmethod +from typing import ClassVar from components.prompts import HYDE_PROMPT, MULTI_QUERY_PROMPT from langchain_core.documents.base import Document @@ -32,9 +33,7 @@ async def retrieve(self, partition: list[str], query: str) -> list[Document]: # Define the Simple Retriever class class BaseRetriever(ABCRetriever): - def __init__( - self, top_k=6, similarity_threshold=0.95, with_surrounding_chunks=True, **kwargs - ): + def __init__(self, top_k=6, similarity_threshold=0.95, with_surrounding_chunks=True, **kwargs): super().__init__(top_k, similarity_threshold, **kwargs) self.top_k = top_k self.similarity_threshold = similarity_threshold @@ -76,12 +75,8 @@ def __init__( if llm is None: raise ValueError("llm must be provided for MultiQueryRetriever") - prompt: ChatPromptTemplate = ChatPromptTemplate.from_template( - MULTI_QUERY_PROMPT - ) - self.generate_queries = ( - prompt | llm | StrOutputParser() | (lambda x: x.split("[SEP]")) - ) + prompt: ChatPromptTemplate = ChatPromptTemplate.from_template(MULTI_QUERY_PROMPT) + self.generate_queries = prompt | llm | StrOutputParser() | (lambda x: x.split("[SEP]")) async def retrieve(self, partition: list[str], query: str) -> list[Document]: db = get_vectordb() @@ -143,7 +138,7 @@ async def retrieve(self, partition: list[str], query: str) -> list[Document]: class RetrieverFactory: - RETRIEVERS = { + RETRIEVERS: ClassVar[dict] = { "single": SingleRetriever, "multiQuery": MultiQueryRetriever, "hyde": HyDeRetriever, diff --git a/openrag/components/utils.py b/openrag/components/utils.py index 76233ee4c..25376e05c 100644 --- a/openrag/components/utils.py +++ b/openrag/components/utils.py @@ -1,5 +1,6 @@ import asyncio import threading +from typing import ClassVar import ray from config import load_config @@ -14,7 +15,7 @@ class SingletonMeta(type): - _instances = {} + _instances: ClassVar[dict] = {} _lock = threading.Lock() # Ensures thread safety def __call__(cls, *args, **kwargs): @@ -83,7 +84,7 @@ def format_context(docs: list[Document], max_context_tokens: int = 4096) -> str: llm = ChatOpenAI(**config.llm) _length_function = llm.get_num_tokens - docs_with_tokens = list(map(lambda d: (_length_function(d.page_content), d), docs)) + docs_with_tokens = list(map(lambda d: (_length_function(d.page_content), d), docs)) # noqa: C417 reduced_docs = [] diff --git a/openrag/models/indexer.py b/openrag/models/indexer.py index 907c86a8c..52aafed1f 100644 --- a/openrag/models/indexer.py +++ b/openrag/models/indexer.py @@ -1,8 +1,7 @@ -from typing import Optional from pydantic import BaseModel class SearchRequest(BaseModel): query: str - top_k: Optional[int] = 5 # default to 5 if not provided + top_k: int | None = 5 # default to 5 if not provided diff --git a/openrag/models/openai.py b/openrag/models/openai.py index c45488570..9b766d4cd 100644 --- a/openrag/models/openai.py +++ b/openrag/models/openai.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, Literal from pydantic import BaseModel, Field @@ -14,14 +14,14 @@ class OpenAIMessage(BaseModel): class OpenAIChatCompletionRequest(BaseModel): """Modèle représentant une requête de complétion chat pour l'API OpenAI.""" - model: Optional[str] = Field(None, description="model name") - messages: List[OpenAIMessage] - temperature: Optional[float] = Field(0.3) - top_p: Optional[float] = Field(1.0) - stream: Optional[bool] = Field(False) - max_tokens: Optional[int] = Field(1024) - logprobs: Optional[int] = Field(None) - metadata: Optional[Dict[str, Any]] = Field( + model: str | None = Field(None, description="model name") + messages: list[OpenAIMessage] + temperature: float | None = Field(0.3) + top_p: float | None = Field(1.0) + stream: bool | None = Field(False) + max_tokens: int | None = Field(1024) + logprobs: int | None = Field(None) + metadata: dict[str, Any] | None = Field( { "use_map_reduce": False, "spoken_style_answer": False, @@ -30,111 +30,21 @@ class OpenAIChatCompletionRequest(BaseModel): ) -class OpenAIChatCompletionChoice(BaseModel): - """Modèle représentant un choix de complétion chat dans l'API OpenAI.""" - - index: int - message: OpenAIMessage - finish_reason: str - - -class OpenAILogprobs(BaseModel): - text_offset: Optional[List] = Field(None) - token_logprobs: Optional[List[float]] = Field(None) - tokens: Optional[List[str]] = Field(None) - top_logprobs: Optional[List] = Field(None) - - -class ChatCompletionTokenLogprob(BaseModel): - token: str - bytes_: list[int] = Field(..., alias="bytes") # Handle Python reserved word - logprob: float - top_logprobs: list # Separate model for recursion - - -class ChoiceLogprobs(BaseModel): - text_offsets: Optional[List[int]] = Field(None) - token_logprobs: List[Union[float, None]] - tokens: List[str] = Field(None) - top_logprobs: List[Union[dict, None]] - - -# Handle forward references -ChatCompletionTokenLogprob.model_rebuild() - - -class OpenAICompletionChoice(BaseModel): - """Modèle représentant un choix de complétion dans l'API OpenAI.""" - - index: int - text: str - logprobs: Optional[ChoiceLogprobs] = Field(None) - finish_reason: str - - class OpenAICompletionRequest(BaseModel): """Legacy OpenAI completion API""" - model: Optional[str] = Field(None, description="model name") + model: str | None = Field(None, description="model name") prompt: str - best_of: Optional[int] = Field(1) - echo: Optional[bool] = Field(False) - frequency_penalty: Optional[float] = Field(0.0) - logit_bias: Optional[dict] = Field(None) - logprobs: Optional[int] = Field(None) - max_tokens: Optional[int] = Field(512) - n: Optional[int] = Field(1) - presence_penalty: Optional[float] = Field(0.0) - seed: Optional[int] = Field(None) - stop: Optional[List[str]] = Field(None) - stream: Optional[bool] = Field(False) - temperature: Optional[float] = Field(0.3) - top_p: Optional[float] = Field(1.0) - - -class OpenAIUsage(BaseModel): - """Modèle représentant les statistiques d'utilisation dans l'API OpenAI.""" - - prompt_tokens: int - completion_tokens: int - total_tokens: int - - -class OpenAICompletion(BaseModel): - """Modèle représentant une réponse de complétion dans l'API OpenAI.""" - - id: str - object: str = "text_completion" - created: int - model: str - choices: List[OpenAICompletionChoice] - usage: OpenAIUsage - - -class OpenAIChatCompletion(BaseModel): - """Modèle représentant une réponse de complétion chat dans l'API OpenAI.""" - - id: str - object: str = "chat.completion" - created: int - model: str - choices: List[OpenAIChatCompletionChoice] - usage: OpenAIUsage - - -class OpenAICompletionChunkChoice(BaseModel): - """Modèle représentant un choix de segment de complétion en streaming dans l'API OpenAI.""" - - index: int - delta: Dict[str, str] - finish_reason: Optional[str] = None - - -class OpenAICompletionChunk(BaseModel): - """Modèle représentant un segment de complétion en streaming dans l'API OpenAI.""" - - id: str - object: str = "chat.completion.chunk" - created: int - model: str - choices: List[OpenAICompletionChunkChoice] + best_of: int | None = Field(1) + echo: bool | None = Field(False) + frequency_penalty: float | None = Field(0.0) + logit_bias: dict | None = Field(None) + logprobs: int | None = Field(None) + max_tokens: int | None = Field(512) + n: int | None = Field(1) + presence_penalty: float | None = Field(0.0) + seed: int | None = Field(None) + stop: list[str] | None = Field(None) + stream: bool | None = Field(False) + temperature: float | None = Field(0.3) + top_p: float | None = Field(1.0) diff --git a/openrag/routers/actors.py b/openrag/routers/actors.py index 4e8cbb2d6..68c025174 100644 --- a/openrag/routers/actors.py +++ b/openrag/routers/actors.py @@ -30,7 +30,9 @@ } -@router.get("/", name="list_ray_actors", +@router.get( + "/", + name="list_ray_actors", description="""List all Ray actors and their current status. **Permissions:** @@ -69,7 +71,9 @@ async def list_ray_actors(): ) -@router.post("/{actor_name}/restart", name="restart_ray_actor", +@router.post( + "/{actor_name}/restart", + name="restart_ray_actor", description="""Restart a specific Ray actor. **Parameters:** @@ -118,7 +122,7 @@ async def restart_actor( logger.exception("Failed to kill actor", actor=actor_name) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Failed to kill actor {actor_name}: {str(e)}", + detail=f"Failed to kill actor {actor_name}: {e!s}", ) try: @@ -138,5 +142,5 @@ async def restart_actor( logger.exception("Failed to restart actor", actor=actor_name) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Failed to restart actor {actor_name}: {str(e)}", + detail=f"Failed to restart actor {actor_name}: {e!s}", ) diff --git a/openrag/routers/extract.py b/openrag/routers/extract.py index be3f39751..5de06514f 100644 --- a/openrag/routers/extract.py +++ b/openrag/routers/extract.py @@ -11,7 +11,8 @@ router = APIRouter() -@router.get("/{extract_id}", +@router.get( + "/{extract_id}", description="""Get a specific document chunk by its ID. **Parameters:** @@ -72,7 +73,7 @@ async def get_extract( log.exception("Failed to retrieve extract.", error=str(e)) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Failed to retrieve extract: {str(e)}", + detail=f"Failed to retrieve extract: {e!s}", ) return JSONResponse( diff --git a/openrag/routers/indexer.py b/openrag/routers/indexer.py index 6b351a216..8f023d549 100644 --- a/openrag/routers/indexer.py +++ b/openrag/routers/indexer.py @@ -1,7 +1,7 @@ import json from datetime import datetime from pathlib import Path -from typing import Any, Optional +from typing import Any import ray from components.indexer.utils.files import sanitize_filename, save_file_to_disk @@ -104,7 +104,7 @@ async def get_supported_types(): **Common Mimetypes:** - `text/plain` - Plain text files -- `text/markdown` - Markdown files +- `text/markdown` - Markdown files - `application/pdf` - PDF documents - `message/rfc822` - Email files @@ -163,18 +163,12 @@ async def add_file( metadata["file_id"] = file_id # Indexing the file - task = indexer.add_file.remote( - path=file_path, metadata=metadata, partition=partition, user=user - ) + task = indexer.add_file.remote(path=file_path, metadata=metadata, partition=partition, user=user) await task_state_manager.set_state.remote(task.task_id().hex(), "QUEUED") await task_state_manager.set_object_ref.remote(task.task_id().hex(), {"ref": task}) return JSONResponse( status_code=status.HTTP_201_CREATED, - content={ - "task_status_url": build_url( - request, "get_task_status", task_id=task.task_id().hex() - ) - }, + content={"task_status_url": build_url(request, "get_task_status", task_id=task.task_id().hex())}, ) @@ -285,19 +279,13 @@ async def put_file( metadata["file_id"] = file_id # Indexing the file - task = indexer.add_file.remote( - path=file_path, metadata=metadata, partition=partition - ) + task = indexer.add_file.remote(path=file_path, metadata=metadata, partition=partition, user=user) await task_state_manager.set_state.remote(task.task_id().hex(), "QUEUED") await task_state_manager.set_object_ref.remote(task.task_id().hex(), {"ref": task}) return JSONResponse( status_code=status.HTTP_202_ACCEPTED, - content={ - "task_status_url": build_url( - request, "get_task_status", task_id=task.task_id().hex() - ) - }, + content={"task_status_url": build_url(request, "get_task_status", task_id=task.task_id().hex())}, ) @@ -322,7 +310,7 @@ async def put_file( async def patch_file( partition: str, file_id: str = Depends(validate_file_id), - metadata: Optional[Any] = Depends(validate_metadata), + metadata: Any | None = Depends(validate_metadata), indexer=Depends(get_indexer), user=Depends(require_partition_editor), user_partitions=Depends(current_user_partitions), @@ -367,7 +355,7 @@ async def patch_file( async def copy_file_between_partitions( partition: str, file_id: str = Depends(validate_file_id), - metadata: Optional[Any] = Depends(validate_metadata), + metadata: Any | None = Depends(validate_metadata), source_partition: str = Form(...), source_file_id: str = Form(...), indexer=Depends(get_indexer), @@ -384,9 +372,7 @@ async def copy_file_between_partitions( metadata["file_id"] = file_id metadata["partition"] = partition - await indexer.copy_file.remote( - file_id=source_file_id, metadata=metadata, partition=source_partition, user=user - ) + await indexer.copy_file.remote(file_id=source_file_id, metadata=metadata, partition=source_partition, user=user) return JSONResponse( status_code=status.HTTP_201_CREATED, content={"message": "File copied successfully."}, @@ -488,15 +474,13 @@ async def get_task_error( **Note:** Logs are returned in chronological order (oldest first). """, ) -async def get_task_logs( - task_id: str, max_lines: int = 100, task_details=Depends(require_task_owner) -): +async def get_task_logs(task_id: str, max_lines: int = 100, task_details=Depends(require_task_owner)): try: if not LOG_FILE.exists(): raise HTTPException(status_code=500, detail="Log file not found.") logs = [] - with open(LOG_FILE, "r", errors="replace") as f: + with open(LOG_FILE, errors="replace") as f: for line in reversed(list(f)): try: record = json.loads(line).get("record", {}) @@ -510,17 +494,13 @@ async def get_task_logs( continue if not logs: - raise HTTPException( - status_code=404, detail=f"No logs found for task '{task_id}'" - ) + raise HTTPException(status_code=404, detail=f"No logs found for task '{task_id}'") - return JSONResponse( - content={"task_id": task_id, "logs": logs[::-1]} - ) # restore order + return JSONResponse(content={"task_id": task_id, "logs": logs[::-1]}) # restore order except HTTPException: raise except Exception as e: - raise HTTPException(status_code=500, detail=f"Failed to fetch logs: {str(e)}") + raise HTTPException(status_code=500, detail=f"Failed to fetch logs: {e!s}") @router.delete( diff --git a/openrag/routers/openai.py b/openrag/routers/openai.py index 8856ddba1..66deab2ac 100644 --- a/openrag/routers/openai.py +++ b/openrag/routers/openai.py @@ -167,7 +167,7 @@ async def openai_chat_completion( status_code=status.HTTP_400_BAD_REQUEST, detail="The last message must be a non-empty user message", ) - + log.debug( "Received chat completion request with messages: {}", truncate(str(request.messages)), @@ -194,7 +194,7 @@ async def openai_chat_completion( log.exception("Chat completion failed.", error=str(e)) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Chat completion failed: {str(e)}", + detail=f"Chat completion failed: {e!s}", ) metadata = __prepare_sources(request2, docs) @@ -245,7 +245,7 @@ async def stream_response(): log.warning("Error while generating answer", error=str(e)) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Error while generating answer: {str(e)}", + detail=f"Error while generating answer: {e!s}", ) @@ -322,7 +322,7 @@ async def openai_completion( log.exception("Completion request failed.", error=str(e)) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Completion failed: {str(e)}", + detail=f"Completion failed: {e!s}", ) metadata = __prepare_sources(request2, docs) @@ -337,5 +337,5 @@ async def openai_completion( log.warning("No response from LLM.", error=str(e)) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"No response from LLM: {str(e)}", + detail=f"No response from LLM: {e!s}", ) diff --git a/openrag/routers/partition.py b/openrag/routers/partition.py index e4af3b38c..17d54643e 100644 --- a/openrag/routers/partition.py +++ b/openrag/routers/partition.py @@ -23,7 +23,8 @@ def _quote_param_value(s: str) -> str: return quote(s, safe="") -@router.get("/", +@router.get( + "/", description="""List all accessible partitions. **Response:** @@ -49,7 +50,8 @@ async def list_existant_partitions( ) -@router.delete("/{partition}", +@router.delete( + "/{partition}", description="""Delete a partition and all its contents. **Parameters:** @@ -75,7 +77,8 @@ async def delete_partition( return Response(status_code=status.HTTP_204_NO_CONTENT) -@router.get("/{partition}", +@router.get( + "/{partition}", description="""List all files in a partition. **Parameters:** @@ -125,7 +128,8 @@ def process_file(file_dict): ) -@router.get("/{partition}/file/{file_id}", +@router.get( + "/{partition}/file/{file_id}", description="""Get details and chunks for a specific file. **Parameters:** @@ -172,7 +176,8 @@ async def get_file( ) -@router.get("/{partition}/chunks", +@router.get( + "/{partition}/chunks", description="""List all document chunks in a partition. **Parameters:** @@ -215,7 +220,8 @@ async def list_all_chunks( return JSONResponse(status_code=status.HTTP_200_OK, content={"chunks": chunks}) -@router.post("/{partition}", +@router.post( + "/{partition}", description="""Create a new partition. **Parameters:** @@ -246,7 +252,8 @@ async def create_partition( return Response(status_code=status.HTTP_201_CREATED) -@router.get("/{partition}/users", +@router.get( + "/{partition}/users", description="""List all users with access to a partition. **Parameters:** @@ -283,7 +290,8 @@ async def list_partition_users( return JSONResponse(status_code=status.HTTP_200_OK, content={"members": members}) -@router.post("/{partition}/users", +@router.post( + "/{partition}/users", description="""Add a user to a partition with a specific role. **Parameters:** @@ -323,7 +331,8 @@ async def add_partition_user( return Response(status_code=status.HTTP_201_CREATED) -@router.delete("/{partition}/users/{user_id}", +@router.delete( + "/{partition}/users/{user_id}", description="""Remove a user from a partition. **Parameters:** @@ -359,7 +368,8 @@ async def remove_partition_user( return Response(status_code=status.HTTP_204_NO_CONTENT) -@router.patch("/{partition}/users/{user_id}", +@router.patch( + "/{partition}/users/{user_id}", description="""Update a user's role in a partition. **Parameters:** diff --git a/openrag/routers/queue.py b/openrag/routers/queue.py index 454c66777..3feeb2246 100644 --- a/openrag/routers/queue.py +++ b/openrag/routers/queue.py @@ -25,7 +25,8 @@ def _format_pool_info(worker_info: dict[str, int]) -> dict[str, int]: } -@router.get("/info", +@router.get( + "/info", description="""Get queue and worker pool information. **Permissions:** @@ -71,7 +72,9 @@ async def get_queue_info( return {"workers": workers_block, "tasks": task_summary} -@router.get("/tasks", name="list_tasks", +@router.get( + "/tasks", + name="list_tasks", description="""List indexing tasks with optional filtering. **Query Parameters:** diff --git a/openrag/routers/search.py b/openrag/routers/search.py index 013e3bba8..896c4afdd 100644 --- a/openrag/routers/search.py +++ b/openrag/routers/search.py @@ -1,4 +1,3 @@ -from typing import List, Optional from fastapi import APIRouter, Depends, Query, Request, status from fastapi.responses import JSONResponse @@ -16,7 +15,8 @@ router = APIRouter() -@router.get("", +@router.get( + "", description="""Perform semantic search across multiple partitions. **Query Parameters:** @@ -46,7 +46,7 @@ ) async def search_multiple_partitions( request: Request, - partitions: Optional[List[str]] = Query( + partitions: list[str] | None = Query( default=["all"], description="List of partitions to search" ), text: str = Query(..., description="Text to search semantically"), @@ -83,7 +83,8 @@ async def search_multiple_partitions( ) -@router.get("/partition/{partition}", +@router.get( + "/partition/{partition}", description="""Perform semantic search within a single partition. **Parameters:** @@ -133,7 +134,8 @@ async def search_one_partition( ) -@router.get("/partition/{partition}/file/{file_id}", +@router.get( + "/partition/{partition}/file/{file_id}", description="""Perform semantic search within a specific file. **Parameters:** diff --git a/openrag/routers/tools.py b/openrag/routers/tools.py index 8847d079c..5f78def72 100644 --- a/openrag/routers/tools.py +++ b/openrag/routers/tools.py @@ -1,6 +1,5 @@ import json from pathlib import Path -from typing import List import ray from components.indexer.utils.files import save_file_to_disk, serialize_file @@ -28,7 +27,7 @@ class ToolInfo(BaseModel): description: str -AVAILABLE_TOOLS: List[ToolInfo] = [ +AVAILABLE_TOOLS: list[ToolInfo] = [ ToolInfo( name="extractText", description="Extract raw text from a file (PDF, Office, audio, etc.)", @@ -63,7 +62,7 @@ def validate_tool(tool: str = Form(...)): @router.get( "/tools", - response_model=List[ToolInfo], + response_model=list[ToolInfo], summary="List available tools", description="""List available tools **Response Format:** diff --git a/openrag/routers/users.py b/openrag/routers/users.py index fc5e48d0a..7f46b74df 100644 --- a/openrag/routers/users.py +++ b/openrag/routers/users.py @@ -1,4 +1,3 @@ -from typing import Optional from fastapi import APIRouter, Depends, Form, HTTPException, Request, Response, status from fastapi.responses import JSONResponse @@ -11,7 +10,8 @@ router = APIRouter() -@router.get("/", +@router.get( + "/", description="""List all users in the system. **Permissions:** @@ -34,7 +34,8 @@ async def list_users(vectordb=Depends(get_vectordb), admin_user=Depends(require_ return JSONResponse(status_code=status.HTTP_200_OK, content={"users": users}) -@router.get("/info", +@router.get( + "/info", description="""Get current authenticated user information. **Authentication:** @@ -56,7 +57,8 @@ async def get_current_user(request: Request): return user -@router.post("/", +@router.post( + "/", description="""Create a new user account. **Parameters:** @@ -79,8 +81,8 @@ async def get_current_user(request: Request): """, ) async def create_user( - display_name: Optional[str] = Form(None), - external_user_id: Optional[str] = Form(None), + display_name: str | None = Form(None), + external_user_id: str | None = Form(None), is_admin: bool = Form(False), vectordb=Depends(get_vectordb), admin_user=Depends(require_admin), @@ -97,7 +99,8 @@ async def create_user( return JSONResponse(status_code=status.HTTP_201_CREATED, content=user) -@router.get("/{user_id}", +@router.get( + "/{user_id}", description="""Get details for a specific user. **Parameters:** @@ -127,7 +130,8 @@ async def get_user( return JSONResponse(status_code=status.HTTP_200_OK, content=user) -@router.delete("/{user_id}", +@router.delete( + "/{user_id}", description="""Delete a user account. **Parameters:** @@ -155,13 +159,15 @@ async def delete_user( """ if user_id == 1: raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, detail="Cannot delete default admin user." + status_code=status.HTTP_400_BAD_REQUEST, + detail="Cannot delete default admin user.", ) await vectordb.delete_user.remote(user_id) return Response(status_code=status.HTTP_204_NO_CONTENT) -@router.post("/{user_id}/regenerate_token", +@router.post( + "/{user_id}/regenerate_token", description="""Regenerate a user's authentication token. **Parameters:** diff --git a/openrag/routers/utils.py b/openrag/routers/utils.py index bed0a0a92..1f23c18f8 100644 --- a/openrag/routers/utils.py +++ b/openrag/routers/utils.py @@ -1,7 +1,7 @@ import json import os from pathlib import Path -from typing import Any, Optional +from typing import Any import consts from config import load_config @@ -197,7 +197,7 @@ async def validate_file_id(file_id: str): return file_id -async def validate_metadata(metadata: Optional[Any] = Form(None)): +async def validate_metadata(metadata: Any | None = Form(None)): try: processed_metadata = metadata or "{}" processed_metadata = json.loads(processed_metadata) diff --git a/openrag/scripts/backup.py b/openrag/scripts/backup.py index c6a43c269..e6f66d918 100644 --- a/openrag/scripts/backup.py +++ b/openrag/scripts/backup.py @@ -1,23 +1,22 @@ #!/usr/bin/env python3 -import sys -import os import json +import os +import sys +from typing import IO, Any -from typing import Dict, Any, Tuple, Optional, List, IO - -from pymilvus import connections, Collection - -from utils.logger import get_logger from components.indexer.vectordb.utils import PartitionFileManager +from pymilvus import Collection, connections +from utils.logger import get_logger def dump_rdb_part( - out_fh: IO[str], - pfm: PartitionFileManager, - partitions: Dict[str, Dict[str, Any]], - logger: Any, - verbose: bool = False): + out_fh: IO[str], + pfm: PartitionFileManager, + partitions: dict[str, dict[str, Any]], + logger: Any, + verbose: bool = False, +): """ Dumps relational DB data into the backup file: - writes only requested partitions @@ -35,39 +34,49 @@ def dump_rdb_part( """ for part_name in sorted(partitions): # Header - out_fh.write('rdb\n') + out_fh.write("rdb\n") if verbose: - logger.info('Writing rdb data') + logger.info("Writing rdb data") # Partition details - out_fh.write(json.dumps({ 'name': part_name, 'created': partitions[part_name]['created_at'] }, ensure_ascii=False, sort_keys=True) + '\n') + out_fh.write( + json.dumps( + {"name": part_name, "created": partitions[part_name]["created_at"]}, + ensure_ascii=False, + sort_keys=True, + ) + + "\n" + ) try: files = pfm.list_partition_files(part_name) except Exception as e: - logger.error(f'Failed while requesting the list of files in partition \'{part_name}\'\n{e}') + logger.error( + f"Failed while requesting the list of files in partition '{part_name}'\n{e}" + ) raise - files['files'].sort(key=lambda v: v['file_id']) + files["files"].sort(key=lambda v: v["file_id"]) - for f in files['files']: - f.pop('partition', None) - out_fh.write(json.dumps(f, ensure_ascii=False, sort_keys=True) + '\n') + for f in files["files"]: + f.pop("partition", None) + out_fh.write(json.dumps(f, ensure_ascii=False, sort_keys=True) + "\n") # Separator - out_fh.write('\n') + out_fh.write("\n") if verbose: - logger.info(f'Partition \'{part_name}\' - {len(files["files"])} files') + logger.info(f"Partition '{part_name}' - {len(files['files'])} files") def dump_vdb_part( - out_fh: IO[str], - collection: Collection, - partitions: Dict[str, Dict[str, Any]], - logger: Any, - batch_size: int = 1024, - verbose: bool = False): + out_fh: IO[str], + collection: Collection, + partitions: dict[str, dict[str, Any]], + logger: Any, + batch_size: int = 1024, + verbose: bool = False, +): """ Dumps vector DB data into the backup file: - writes one chunk per line @@ -88,48 +97,47 @@ def dump_vdb_part( try: collection.load() except Exception as e: - logger.error(f'Failed while loading Milvus collection: {e}') + logger.error(f"Failed while loading Milvus collection: {e}") raise try: iterator = collection.query_iterator( - batch_size=batch_size, # size of each batch - output_fields=["*"] # all fields + batch_size=batch_size, # size of each batch + output_fields=["*"], # all fields ) except Exception as e: - logger.error(f'Failed while trying to obtain Milvus collection iterator: {e}') + logger.error(f"Failed while trying to obtain Milvus collection iterator: {e}") raise - out_fh.write('vdb\n') + out_fh.write("vdb\n") if verbose: - logger.info('Writing vdb data') + logger.info("Writing vdb data") cnt = 0 while True: try: batch = iterator.next() except Exception as e: - logger.error(f'iterator.next() failed with: {e}') + logger.error(f"iterator.next() failed with: {e}") raise if not batch: # no more data break for entity in batch: - if entity['partition'] in partitions: - entity.pop('_id', None) - out_fh.write(json.dumps(entity, ensure_ascii=False, sort_keys=True) + '\n') + if entity["partition"] in partitions: + entity.pop("_id", None) + out_fh.write( + json.dumps(entity, ensure_ascii=False, sort_keys=True) + "\n" + ) if verbose: cnt += 1 if verbose: - logger.info(f'{cnt} chunks written') + logger.info(f"{cnt} chunks written") -def open_output_file( - file_name: str, - logger: Any - ) -> IO[str]: +def open_output_file(file_name: str, logger: Any) -> IO[str]: """ Opens output file for writing @@ -140,20 +148,23 @@ def open_output_file( Returns: file object: Opened file handle in text mode. """ - if file_name in [ '-' ]: + if file_name in ["-"]: return sys.stdout if os.path.isfile(file_name): - raise Exception(f'File \'{file_name}\' already exists.') + raise Exception(f"File '{file_name}' already exists.") try: - if file_name.endswith('.xz'): + if file_name.endswith(".xz"): import lzma - return lzma.open(file_name, 'wt', encoding='utf-8', preset=9 | lzma.PRESET_EXTREME) + + return lzma.open( + file_name, "wt", encoding="utf-8", preset=9 | lzma.PRESET_EXTREME + ) else: - return open(file_name, 'wt', encoding='utf-8') + return open(file_name, "w", encoding="utf-8") except Exception as e: - logger.error(f'Failed while opening file \'{file_name}\' for writing:\n' + str(e)) + logger.error(f"Failed while opening file '{file_name}' for writing:\n" + str(e)) raise @@ -172,6 +183,7 @@ def main(): Returns: int: Exit code (0 on success, non-zero on failure). """ + def load_openrag_config(logger): """ Loads OpenRAG configuration. @@ -189,19 +201,31 @@ def load_openrag_config(logger): try: config = load_config() except Exception as e: - logger.error(f'Failed while trying to obtain OpenRAG config: {e}') + logger.error(f"Failed while trying to obtain OpenRAG config: {e}") raise - return config['rdb'], config['vectordb'] - + return config["rdb"], config["vectordb"] # Arguments and configs import argparse - parser = argparse.ArgumentParser(description='OpenRAG backup tool') - parser.add_argument('-i', '--include-only', nargs='*', help='Include only listed partitions') - parser.add_argument('-o', '--output', required=True, help='Output file name (- for STDOUT)') - parser.add_argument('-b', '--batch-size', default=1024, type=int, help='Batch size used to iterate Milvus') - parser.add_argument('-v', '--verbose', default=False, action='store_true', help='Be verbose') + + parser = argparse.ArgumentParser(description="OpenRAG backup tool") + parser.add_argument( + "-i", "--include-only", nargs="*", help="Include only listed partitions" + ) + parser.add_argument( + "-o", "--output", required=True, help="Output file name (- for STDOUT)" + ) + parser.add_argument( + "-b", + "--batch-size", + default=1024, + type=int, + help="Batch size used to iterate Milvus", + ) + parser.add_argument( + "-v", "--verbose", default=False, action="store_true", help="Be verbose" + ) args = parser.parse_args() @@ -210,8 +234,9 @@ def load_openrag_config(logger): rdb, vdb = load_openrag_config(logger) if args.verbose: - logger.info(f'rdb @ {rdb["host"]}:{rdb["port"]} | vdb @ {vdb["host"]}:{vdb["port"]} | collection: {vdb["collection_name"]}') - + logger.info( + f"rdb @ {rdb['host']}:{rdb['port']} | vdb @ {vdb['host']}:{vdb['port']} | collection: {vdb['collection_name']}" + ) # List existing partitions try: @@ -220,9 +245,13 @@ def load_openrag_config(logger): logger=logger, ) - existing_partitions = { item['partition']: item for item in pfm.list_partitions() } + existing_partitions = { + item["partition"]: item for item in pfm.list_partitions() + } except Exception as e: - logger.error(f'Failed while accessing PartitionFileManager at {rdb["host"]}:{rdb["port"]}\n{e}') + logger.error( + f"Failed while accessing PartitionFileManager at {rdb['host']}:{rdb['port']}\n{e}" + ) raise if args.include_only: @@ -237,45 +266,50 @@ def load_openrag_config(logger): partitions = existing_partitions if 0 == len(partitions): - logger.error(f'No partitions meet given conditions.') + logger.error("No partitions meet given conditions.") return 1 if args.verbose: - partitions_str = ', '.join(partitions) - logger.info(f'partitions: {partitions_str}') - + partitions_str = ", ".join(partitions) + logger.info(f"partitions: {partitions_str}") # Connect to Milvus try: - connections.connect("default", host=vdb['host'], port=vdb['port']) + connections.connect("default", host=vdb["host"], port=vdb["port"]) except Exception as e: - logger.error(f'Can\'t connect to Milvus at {vdb["host"]}:{vdb["port"]}\n{e}') + logger.error(f"Can't connect to Milvus at {vdb['host']}:{vdb['port']}\n{e}") raise try: - vdb_collection = Collection(vdb['collection_name']) + vdb_collection = Collection(vdb["collection_name"]) except Exception as e: - logger.error(f'Can\'t access Milvus collection {vdb["collection_name"]} at {vdb["host"]}:{vdb["port"]}\n{e}') + logger.error( + f"Can't access Milvus collection {vdb['collection_name']} at {vdb['host']}:{vdb['port']}\n{e}" + ) raise - try: with open_output_file(args.output, logger) as out_fh: # Dump data from RDB (one line per document) dump_rdb_part(out_fh, pfm, partitions, logger, args.verbose) # Dump data from VDB (one line per chunk) - dump_vdb_part(out_fh, vdb_collection, partitions, logger, args.batch_size, args.verbose) + dump_vdb_part( + out_fh, + vdb_collection, + partitions, + logger, + args.batch_size, + args.verbose, + ) out_fh.flush() except Exception as e: - logger.exception(f'ERROR: ' + str(e)) + logger.exception("ERROR: " + str(e)) return 1 - return 0 -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) - diff --git a/openrag/scripts/embed.py b/openrag/scripts/embed.py index 166e8bf78..b79dfee3f 100644 --- a/openrag/scripts/embed.py +++ b/openrag/scripts/embed.py @@ -1,12 +1,11 @@ #!/usr/bin/env python3 -import sys -import os import json -import time import logging - -from typing import IO, Any, Dict, List, Optional, Set, Tuple +import os +import sys +import time +from typing import IO, Any class Embedder: @@ -31,7 +30,8 @@ def __init__( model_name: str, batch_size: int, logger: Any, - verbose: bool = False): + verbose: bool = False, + ): """ Initialize an Embedder instance. @@ -54,14 +54,13 @@ def __init__( self.max_batch_size = 1024 * 4 self.curr_batch_size = 1 - self.durations = { self.curr_batch_size: [] } + self.durations = {self.curr_batch_size: []} self.max_duration_items = 12 self.num_outliers = 2 if self.batch_size <= 0: self.batch_size = None # We'll guess the best value - def get_batch_size(self) -> int: """ Get the current batch size. @@ -75,15 +74,14 @@ def get_batch_size(self) -> int: if self.curr_batch_size >= self.max_batch_size: # We are ready to make a choice self.calc_batch_size() - assert(self.batch_size is not None and self.batch_size > 0) + assert self.batch_size is not None and self.batch_size > 0 return self.batch_size self.curr_batch_size *= 2 return self.curr_batch_size - assert(self.batch_size is not None and self.batch_size > 0) + assert self.batch_size is not None and self.batch_size > 0 return self.batch_size - def calc_batch_size(self) -> None: """ Compute the optimal batch size based on recorded durations. @@ -94,22 +92,21 @@ def calc_batch_size(self) -> None: """ data = [] for k, v in self.durations.items(): - assert(len(v) >= self.max_duration_items) - assert(int == type(k)) + assert len(v) >= self.max_duration_items + assert isinstance(k, int) v.sort() # remove outliers - v = v[self.num_outliers:-self.num_outliers] + v = v[self.num_outliers : -self.num_outliers] mean = float(sum(v)) / len(v) - data.append({ 'batch_size': k, 'time_per_item': mean / k }) + data.append({"batch_size": k, "time_per_item": mean / k}) if self.verbose: - self.logger.info(f'calc_batch_size: batch_size={k} time_per_item={mean/k:.4f} mean={mean:.4f}\n') - - data.sort(key=lambda item: item['time_per_item']) - self.batch_size = data[0]['batch_size'] + self.logger.info(f"calc_batch_size: batch_size={k} time_per_item={mean / k:.4f} mean={mean:.4f}\n") + data.sort(key=lambda item: item["time_per_item"]) + self.batch_size = data[0]["batch_size"] def embed(self, data: list) -> list: """ @@ -123,30 +120,28 @@ def embed(self, data: list) -> list: """ before = time.time() - response = self.client.embeddings.create( - model=self.model_name, - input=data - ) + response = self.client.embeddings.create(model=self.model_name, input=data) elapsed = time.time() - before - l = len(data) - if l not in self.durations: - self.durations[l] = [] - elif len(self.durations[l]) >= self.max_duration_items: - self.durations[l].pop(0) - self.durations[l].append(elapsed) + batch_size = len(data) + if batch_size not in self.durations: + self.durations[batch_size] = [] + elif len(self.durations[batch_size]) >= self.max_duration_items: + self.durations[batch_size].pop(0) + self.durations[batch_size].append(elapsed) - return [ item.embedding for item in response.data ] + return [item.embedding for item in response.data] def call_embedder_and_save( - embedder: Embedder, - ofh: IO[str], - batch: list, - text_field_name: str, - vector_field_name: str, - logger: Any, - verbose: bool = False): + embedder: Embedder, + ofh: IO[str], + batch: list, + text_field_name: str, + vector_field_name: str, + logger: Any, + verbose: bool = False, +): """ Call the embedder on a batch of items and save the results. @@ -160,22 +155,22 @@ def call_embedder_and_save( verbose: If True, logs extra info. """ - embd = embedder.embed([ item[text_field_name] for item in batch ]) + embd = embedder.embed([item[text_field_name] for item in batch]) for i in range(len(batch)): batch[i][vector_field_name] = embd[i] - ofh.write(json.dumps(batch[i], ensure_ascii=False, sort_keys=True) + '\n') + ofh.write(json.dumps(batch[i], ensure_ascii=False, sort_keys=True) + "\n") def read_vdb_section( - ifh: IO[str], - ofh: IO[str], - embedder: Any, - logger: Any, - text_field_name: str = 'text', - vector_field_name: str = 'vector', - verbose: bool = False, - ) -> None: + ifh: IO[str], + ofh: IO[str], + embedder: Any, + logger: Any, + text_field_name: str = "text", + vector_field_name: str = "vector", + verbose: bool = False, +) -> None: """ Reads chunks, deletes old vectors and creates new ones with given embedder. @@ -189,25 +184,32 @@ def read_vdb_section( verbose: If True, logs additional info. """ if verbose: - logger.info(f'Read vdb section') + logger.info("Read vdb section") batch_size = embedder.get_batch_size() batch = [] - cnt = 0 for line in ifh: # End of section if 0 == len(line): break if len(batch) >= batch_size: - call_embedder_and_save(embedder, ofh, batch, text_field_name, vector_field_name, logger, verbose) + call_embedder_and_save( + embedder, + ofh, + batch, + text_field_name, + vector_field_name, + logger, + verbose, + ) batch = [] batch_size = embedder.get_batch_size() chunk = json.loads(line) - chunk.pop('_id', None) + chunk.pop("_id", None) chunk.pop(vector_field_name, None) batch.append(chunk) @@ -215,10 +217,7 @@ def read_vdb_section( call_embedder_and_save(embedder, ofh, batch, text_field_name, vector_field_name, logger, verbose) -def open_input_file( - file_name: str, - logger: Any - ) -> IO[str]: +def open_input_file(file_name: str, logger: Any) -> IO[str]: """ Opens a input file for reading, with support for plain text and LZMA-compressed (.xz) files. @@ -229,24 +228,22 @@ def open_input_file( Returns: file object: Opened file handle in text mode. """ - if file_name in [ '-' ]: + if file_name in ["-"]: return sys.stdin try: - if file_name.endswith('.xz'): + if file_name.endswith(".xz"): import lzma - return lzma.open(file_name, 'rt', encoding='utf-8') + + return lzma.open(file_name, "rt", encoding="utf-8") else: - return open(file_name, 'rt', encoding='utf-8') + return open(file_name, encoding="utf-8") except Exception as e: - logger.error(f'Failed while opening file \'{file_name}\' for reading:\n' + str(e)) + logger.error(f"Failed while opening file '{file_name}' for reading:\n" + str(e)) raise -def open_output_file( - file_name: str, - logger: Any - ) -> IO[str]: +def open_output_file(file_name: str, logger: Any) -> IO[str]: """ Opens output file for writing @@ -257,28 +254,25 @@ def open_output_file( Returns: file object: Opened file handle in text mode. """ - if file_name in [ '-' ]: + if file_name in ["-"]: return sys.stdout if os.path.isfile(file_name): - raise Exception(f'File \'{file_name}\' already exists.') + raise Exception(f"File '{file_name}' already exists.") try: - if file_name.endswith('.xz'): + if file_name.endswith(".xz"): import lzma - return lzma.open(file_name, 'wt', encoding='utf-8', preset=9 | lzma.PRESET_EXTREME) + + return lzma.open(file_name, "wt", encoding="utf-8", preset=9 | lzma.PRESET_EXTREME) else: - return open(file_name, 'wt', encoding='utf-8') + return open(file_name, "w", encoding="utf-8") except Exception as e: - logger.error(f'Failed while opening file \'{file_name}\' for writing:\n' + str(e)) + logger.error(f"Failed while opening file '{file_name}' for writing:\n" + str(e)) raise -def close_file( - file_handle: IO[str], - file_name: str, - logger: Any - ) -> None: +def close_file(file_handle: IO[str], file_name: str, logger: Any) -> None: """ Depending of the file name closes the file handle or does nothing @@ -290,13 +284,13 @@ def close_file( Returns: Nothing """ - if file_name in [ '-' ] or file_handle is None: + if file_name in ["-"] or file_handle is None: return try: file_handle.close() except Exception as e: - logger.exception(f'Failed to close file \'{file_name}\': ' + str(e)) + logger.exception(f"Failed to close file '{file_name}': " + str(e)) raise @@ -312,27 +306,38 @@ def main(): int: Exit code (0 on success, non-zero on failure). """ - # Arguments and configs import argparse - parser = argparse.ArgumentParser(description='OpenRAG embed tool') - parser.add_argument('-b', '--batch-size', default=0, type=int, help='Batch size (0 - guess optimal batch size)') - parser.add_argument('-v', '--verbose', default=False, action='store_true', help='Be verbose') - parser.add_argument('-i', '--input', default='-', type=str, help='Input file name (\'-\' for STDIN)') - parser.add_argument('-o', '--output', default='-', type=str, help='Output file name (\'-\' for STDOUT)') - parser.add_argument('-u', '--url', type=str, help='URL to embedder OpenAI compatible endpoint') - parser.add_argument('-k', '--key', type=str, help='Secret key to access embedder') - parser.add_argument('-m', '--model', required=True, type=str, help='Model name') + + parser = argparse.ArgumentParser(description="OpenRAG embed tool") + parser.add_argument( + "-b", + "--batch-size", + default=0, + type=int, + help="Batch size (0 - guess optimal batch size)", + ) + parser.add_argument("-v", "--verbose", default=False, action="store_true", help="Be verbose") + parser.add_argument("-i", "--input", default="-", type=str, help="Input file name ('-' for STDIN)") + parser.add_argument( + "-o", + "--output", + default="-", + type=str, + help="Output file name ('-' for STDOUT)", + ) + parser.add_argument("-u", "--url", type=str, help="URL to embedder OpenAI compatible endpoint") + parser.add_argument("-k", "--key", type=str, help="Secret key to access embedder") + parser.add_argument("-m", "--model", required=True, type=str, help="Model name") args = parser.parse_args() logger = logging.getLogger(__name__) - try: embedder = Embedder(args.url, args.key, args.model, args.batch_size, logger, args.verbose) except Exception as e: - logger.exception(f'Failed while trying to create embedder: ' + str(e)) + logger.exception("Failed while trying to create embedder: " + str(e)) raise try: @@ -343,17 +348,16 @@ def main(): for line in ifh: ofh.write(line) - if line.strip() in [ 'vdb' ]: - read_vdb_section(ifh, ofh, embedder, logger, 'text', 'vector', args.verbose) + if line.strip() in ["vdb"]: + read_vdb_section(ifh, ofh, embedder, logger, "text", "vector", args.verbose) except Exception as e: - logger.exception(f'Error: ' + str(e)) + logger.exception("Error: " + str(e)) raise finally: close_file(ifh, args.input, logger) close_file(ofh, args.output, logger) -if __name__ == '__main__': +if __name__ == "__main__": sys.exit(main()) - diff --git a/openrag/scripts/filter-logs.py b/openrag/scripts/filter-logs.py index 4a1b99dc8..1bc5d2b1b 100644 --- a/openrag/scripts/filter-logs.py +++ b/openrag/scripts/filter-logs.py @@ -2,8 +2,7 @@ import argparse import json import re -from datetime import datetime, timezone -from typing import Optional +from datetime import UTC, datetime, timezone # Expected format at start of record["text"]: # "2025-12-11 09:00:42.538 | INFO ..." @@ -11,7 +10,7 @@ r"^(?P\d{4}-\d{2}-\d{2})\s+(?P