From 224de9eaa15cdd73715f577571de30208821a699 Mon Sep 17 00:00:00 2001 From: Aleksander Obuchowski Date: Thu, 4 Sep 2025 16:13:05 +0200 Subject: [PATCH] ci: add lint and test workflows with uv setup feat(rocketrag): initialize RAG instance in constructor to avoid recreation test: update integration tests for streaming response handling fix(vectors): add PYTORCH_ENABLE_MPS_FALLBACK environment variable --- .github/actions/setup-uv/action.yml | 33 +++++++++++++++++++++++++++ .github/workflows/lint.yml | 26 +++++++++++++++++++++ .github/workflows/test.yml | 35 +++++++++++++++++++++++++++++ rocketrag/rocketrag.py | 8 +++---- rocketrag/vectors.py | 3 +++ tests/test_integration.py | 15 +++++++++---- 6 files changed, 112 insertions(+), 8 deletions(-) create mode 100644 .github/actions/setup-uv/action.yml create mode 100644 .github/workflows/lint.yml create mode 100644 .github/workflows/test.yml diff --git a/.github/actions/setup-uv/action.yml b/.github/actions/setup-uv/action.yml new file mode 100644 index 0000000..9da3ac1 --- /dev/null +++ b/.github/actions/setup-uv/action.yml @@ -0,0 +1,33 @@ +name: 'Setup uv' +description: 'Setup uv with Python and install dependencies' +inputs: + python-version: + description: 'Python version to use' + required: false + default: '3.12' + install-dev: + description: 'Install dev dependencies' + required: false + default: 'true' + +runs: + using: 'composite' + steps: + - name: Install uv + uses: astral-sh/setup-uv@v4 + with: + enable-cache: true + cache-dependency-glob: "uv.lock" + + - name: Set up Python + run: uv python install ${{ inputs.python-version }} + shell: bash + + - name: Install dependencies + run: | + if [ "${{ inputs.install-dev }}" = "true" ]; then + uv sync --all-extras --dev + else + uv sync --all-extras + fi + shell: bash \ No newline at end of file diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 0000000..40d6b30 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,26 @@ +name: Lint + +on: + push: + branches: [ "*" ] + workflow_dispatch: + +jobs: + lint: + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Setup uv + uses: ./.github/actions/setup-uv + with: + python-version: '3.12' + install-dev: 'true' + + - name: Run ruff check + run: uv run ruff check . + + - name: Run ruff format check + run: uv run ruff format --check . \ No newline at end of file diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..91b4ad4 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,35 @@ +name: Test + +on: + pull_request: + branches: [ "main", "dev" ] + workflow_dispatch: + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.12'] + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Setup uv + uses: ./.github/actions/setup-uv + with: + python-version: ${{ matrix.python-version }} + install-dev: 'true' + + - name: Run tests + run: uv run pytest -v --cov=rocketrag --cov-report=term-missing + + - name: Run type checking (if available) + run: | + if uv run python -c "import mypy" 2>/dev/null; then + uv run mypy rocketrag + else + echo "Type checker not available, skipping" + fi + continue-on-error: true \ No newline at end of file diff --git a/rocketrag/rocketrag.py b/rocketrag/rocketrag.py index b62c7db..9746b2d 100644 --- a/rocketrag/rocketrag.py +++ b/rocketrag/rocketrag.py @@ -79,6 +79,8 @@ def __init__( self.metadata, ) + self.rag = RAG(self.db, self.llm) + def prepare(self, recreate: bool = False): if self.loader is None: raise ValueError("Loader is not defined.") @@ -100,13 +102,11 @@ def stream_llm(self, messages: list[dict]) -> str: @ensure_llm_loaded def ask(self, question: str) -> tuple[str, list[SearchResult]]: # TODO: Should print the error on vectorizer mismach - rag = RAG(self.db, self.llm) - stream, sources = rag.run(question) + stream, sources = self.rag.run(question) return stream, sources @ensure_llm_loaded def stream_ask(self, question: str) -> tuple[str, list[SearchResult]]: # TODO: Should print the error on vectorizer mismach - rag = RAG(self.db, self.llm) - stream, sources = rag.stream(question) + stream, sources = self.rag.stream(question) return stream, sources diff --git a/rocketrag/vectors.py b/rocketrag/vectors.py index be6ad30..ffd452b 100644 --- a/rocketrag/vectors.py +++ b/rocketrag/vectors.py @@ -1,3 +1,6 @@ +import os + +os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" from sentence_transformers import SentenceTransformer from .base import BaseVectorizer diff --git a/tests/test_integration.py b/tests/test_integration.py index c192435..5f377dd 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -115,11 +115,18 @@ def test_stream_ask_question(self, integration_rocketrag): # Ask a question using streaming question = "What are the main topics covered?" - response, sources = integration_rocketrag.stream_ask(question) + stream, sources = integration_rocketrag.stream_ask(question) + + # Consume the stream to get the complete response + full_response = "" + for output in stream: + delta = output["choices"][0]["delta"] + if "content" in delta: + full_response += delta["content"] # Verify response (streaming should still return complete response) - assert isinstance(response, str) - assert len(response) > 0 + assert isinstance(full_response, str) + assert len(full_response) > 0 # Verify sources assert isinstance(sources, list) @@ -204,7 +211,7 @@ def test_recreate_database(self, integration_rocketrag): """Test recreating the database.""" # Prepare the database integration_rocketrag.prepare() - initial_count = integration_rocketrag.db.count_records() + initial_count = integration_rocketrag.db.get_total_count() assert initial_count > 0 # Recreate the database