diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000000..71177351f1 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,24 @@ +## Summary + + + +## Type of change + +- [ ] Bug fix +- [ ] New feature (reader / embedder / generator / chunker) +- [ ] Documentation update +- [ ] Refactor / cleanup +- [ ] Other + +## Testing + + + +- [ ] Tested locally against a running Verba instance +- [ ] Added / updated tests in `goldenverba/tests/` + +## Checklist + +- [ ] My code follows the existing patterns in the codebase +- [ ] I have added environment variable docs to `.env.example` if applicable +- [ ] I have updated `CHANGELOG.md` under the relevant version section diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 850060d8be..e76633abdb 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -5,26 +5,39 @@ updates: schedule: interval: weekly day: monday - open-pull-requests-limit: 10 + open-pull-requests-limit: 3 labels: - dependencies - python + groups: + python-deps: + patterns: + - "*" - package-ecosystem: npm directory: "/frontend" schedule: interval: weekly day: monday - open-pull-requests-limit: 10 + open-pull-requests-limit: 3 labels: - dependencies - javascript + groups: + npm-deps: + patterns: + - "*" - package-ecosystem: github-actions directory: "/" schedule: interval: weekly day: monday + open-pull-requests-limit: 1 labels: - dependencies - github-actions + groups: + github-actions: + patterns: + - "*" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000000..6cd607ffee --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,58 @@ +name: CI + +on: + push: + branches: ["main", "v3"] + pull_request: + branches: ["main", "v3"] + +jobs: + backend-test: + name: Backend (Python ${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.11", "3.12"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + + - name: Install dependencies + run: | + pip install --upgrade pip + pip install ruff + pip install '.[dev]' + + - name: Lint with ruff + run: ruff check goldenverba/ + + - name: Run tests + run: pytest goldenverba/tests/ -v --tb=short + + frontend-lint: + name: Frontend (Node 20) + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Node.js + uses: actions/setup-node@v4 + with: + node-version: "20" + cache: npm + cache-dependency-path: frontend/package-lock.json + + - name: Install dependencies + working-directory: frontend + run: npm ci + + - name: Lint + working-directory: frontend + run: npm run lint diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml index 10055aac91..94b6077ba2 100644 --- a/.github/workflows/docker-image.yml +++ b/.github/workflows/docker-image.yml @@ -22,10 +22,12 @@ jobs: password: ${{secrets.DOCKER_PASSWORD}} - name: Build and push - uses: docker/build-push-action@v4 + uses: docker/build-push-action@v6 with: context: . file: ./Dockerfile push: true tags: semitechnologies/verba:latest platforms: linux/amd64,linux/arm64 + cache-from: type=gha + cache-to: type=gha,mode=max diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000000..5ba23e10fe --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,67 @@ +name: Release + +on: + push: + tags: + - "v*.*.*" + +jobs: + pypi-publish: + name: Publish to PyPI + runs-on: ubuntu-latest + environment: + name: pypi + url: https://pypi.org/project/goldenverba/ + permissions: + id-token: write # Required for trusted publishing + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install build tools + run: pip install build + + - name: Build package + run: python -m build + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + + docker-push: + name: Build and push Docker image + runs-on: ubuntu-latest + needs: pypi-publish + + steps: + - uses: actions/checkout@v4 + + - name: Extract version tag + id: meta + run: echo "tag=${GITHUB_REF_NAME}" >> $GITHUB_OUTPUT + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to DockerHub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKER_USERNAME }} + password: ${{ secrets.DOCKER_PASSWORD }} + + - name: Build and push versioned + latest tags + uses: docker/build-push-action@v6 + with: + context: . + file: ./Dockerfile + push: true + tags: | + semitechnologies/verba:${{ steps.meta.outputs.tag }} + semitechnologies/verba:latest + platforms: linux/amd64,linux/arm64 + cache-from: type=gha + cache-to: type=gha,mode=max diff --git a/.gitignore b/.gitignore index e9e47ddcd9..baf19b1532 100644 --- a/.gitignore +++ b/.gitignore @@ -1,24 +1,39 @@ +# Environment .env .env* + +# Python __pycache__ -.DS_Store -.pytest_cache .python-version *.egg-info +.ruff_cache +.pytest_cache + +# Virtual environments venv .venv venv* + +# Build artifacts dist build -~ + +# Local data and config +data/ +verba_config.json +*_secrets.json +cache.txt + +# OS +.DS_Store + +# Tools .local .cache .verba .vscode -verba_config.json +ollama text-generation-inference test.py -cache.txt -.ruff_cache -*_secrets.json -ollama + +todo.md \ No newline at end of file diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000000..f59dd19c3c --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,27 @@ +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + - id: check-json + - id: check-merge-conflict + - id: check-added-large-files + args: ["--maxkb=1000"] + - id: detect-private-key + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.11.6 + hooks: + - id: ruff + args: ["--fix"] + - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-prettier + rev: v4.0.0-alpha.8 + hooks: + - id: prettier + types_or: [javascript, jsx, ts, tsx, json, css, markdown] + files: ^frontend/ + exclude: frontend/node_modules/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 8bd7e5d2a1..52ca2c1a74 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,64 @@ All notable changes to this project will be documented in this file. +## [3.0.0] - Unreleased + +## Added + +- **WhisperReader** — local audio/video transcription via `faster-whisper`. Replaces the AssemblyAI reader with a zero-cost, offline alternative. Supports 35+ formats (`.mp3`, `.wav`, `.m4a`, `.flac`, `.ogg`, `.mp4`, `.mov`, `.webm`, …). No API key required; model is downloaded automatically on first use. Configurable model size (`tiny` → `large-v3`) and compute device (`cpu` / `cuda` / `auto`). + +## Removed + +- **FirecrawlReader** — removed due to high maintenance surface (async job polling, versioned API endpoints, paid usage-based pricing). HTMLReader covers the primary use case (static site scraping). Users needing JavaScript rendering should run a local Playwright/Puppeteer setup. +- **UpstageDocumentParse reader** — removed as redundant with UnstructuredAPI for complex PDF parsing. Upstage embedding and generation models are unaffected and remain available. +- **AssemblyAI reader** — replaced by the new local WhisperReader. Removes a paid-per-minute external dependency while providing the same audio/video transcription capability. + +## Added + +- **DeepSeek generator** with reasoning model (R1) support — supports `deepseek-chat` and `deepseek-reasoner` via DeepSeek's OpenAI-compatible API. R1 thinking process shown in a collapsible section with a "Show Reasoning" toggle. Dynamic model discovery at startup. Env vars: `DEEPSEEK_API_KEY`, `DEEPSEEK_BASE_URL`, `DEEPSEEK_MODEL` (https://github.com/weaviate/Verba/pull/395) +- **LM Studio integration** — `LMStudioEmbedder` and `LMStudioGenerator` for running fully local models via LM Studio's OpenAI-compatible API (`http://localhost:1234/v1` by default). No API key required. Env vars: `LMSTUDIO_BASE_URL`, `LMSTUDIO_API_KEY`, `LMSTUDIO_MODEL`, `LMSTUDIO_EMBEDDER_MODEL` (https://github.com/weaviate/Verba/pull/391) +- **Test suite** — 71 tests covering `BatchManager` (including TTL eviction), `LoggerManager`, all three chunkers (`TokenChunker`, `SentenceChunker`, `MarkdownChunker`), and the `Document` class +- **BACKEND.md** — developer guide covering architecture, component plugin system, step-by-step guides for adding new Generators/Embedders/Readers/Chunkers, WebSocket protocol, config system, and local dev setup + +## Fixed + +- Chunk deserialization from JSON: `doc_uuid` was stored as a tuple due to a stray trailing comma; `title`, `labels`, and `pca` were silently dropped on round-trip. Includes a full serialization round-trip test. (https://github.com/weaviate/Verba/pull/398) +- CORS misconfiguration: `allow_credentials=True` with `allow_origins=["*"]` is rejected by browsers per the CORS spec. Access control is enforced by the existing custom same-origin middleware. (https://github.com/weaviate/Verba/issues/393) +- **Mutable default argument** in `import_document()` — `LoggerManager()` was instantiated once at definition time and shared across concurrent imports, causing them to clobber each other's WebSocket reference +- **Lock creation race** in `ClientManager.get_or_create_lock()` — two coroutines could both see an absent key and create duplicate locks; fixed with `setdefault()` +- **Dict mutation during iteration** in `ClientManager.clean_up()` and `disconnect()` — iterating over a live dict while deleting entries raises `RuntimeError`; fixed by iterating over a snapshot +- **Exception object sent to WebSocket** — `send_json({"message": e})` where `e` is an `Exception` object fails JSON serialization; changed to `str(e)` +- **`msg.good()` fired after exceptions** — success log was outside the `try` block and ran unconditionally even when the streaming failed +- **`asyncio.create_task()` immediately awaited** — pattern `task = create_task(fn()); result = await task` is equivalent to `await fn()` but with extra overhead; simplified to direct `await` +- **O(n²) string concatenation** in streaming loop — `full_text += chunk["message"]` in a hot loop replaced with list accumulation and `"".join()` at the end +- **`SentenceTransformer` model reloaded on every call** — `SentenceTransformer(model_name)` was called inside `vectorize()`, loading ~300 MB from disk on every embedding request; model is now cached by name in `_model_cache` +- **`model.encode()` blocking the event loop** — synchronous CPU-bound call now wrapped with `asyncio.to_thread()` +- **Streaming timeout `None`** on all generators — a hung upstream API would hold a WebSocket connection open forever; all generators now use `httpx.Timeout(connect=10, read=300)` or `aiohttp.ClientTimeout(connect=10, total=300)` +- **Debug `print()` statements** in `util.py` `pca()` — four statements printing raw matrices to stdout in production removed +- **Bare `except:` clauses** — `except:` in `GeminiGenerator.py` (import guard) and `document.py` (language detection) changed to `except ImportError` and `except Exception` respectively +- **Deprecated `asyncio==3.4.3`** in `setup.py` — the PyPI `asyncio` package conflicts with the stdlib module present since Python 3.4; entry removed + +## Changed + +- **Payload size limits** — `GeneratePayload` and `QueryPayload` now enforce Pydantic `max_length` validators: query ≤ 50,000 chars, context ≤ 500,000 chars, conversation ≤ 100 items +- **`BatchManager` TTL eviction** — abandoned incomplete uploads are now evicted after 5 minutes (previously leaked memory indefinitely) +- Automated PyPI publishing via GitHub Actions on `v*.*.*` tag push using trusted publishing — no stored API token needed (replaces manual `pypi_commands.sh`) +- Docker image now also tagged with version (e.g. `semitechnologies/verba:v3.0.0`) in addition to `:latest` +- Upgraded Docker GitHub Actions to `build-push-action@v6` with GHA build cache for faster builds +- `httpx` now explicitly declared in `setup.py` (was already used by multiple generators but only present as a transitive dependency) + +## Infrastructure + +- CI workflow (`ci.yml`): runs pytest on Python 3.11/3.12, ruff linting, and ESLint on every PR targeting `main` or `v3` +- `ruff.toml`: Python linting and formatting config (replaces Black) +- `.pre-commit-config.yaml`: ruff, ruff-format, prettier for frontend, and file hygiene hooks +- Dependabot: automated weekly dependency updates for pip, npm, and GitHub Actions (grouped into one PR per ecosystem) +- `SECURITY.md`: responsible disclosure policy via GitHub private vulnerability reporting +- `CODE_OF_CONDUCT.md`: Contributor Covenant 2.1 +- `.github/PULL_REQUEST_TEMPLATE.md`: PR checklist for contributors + +--- + ## [2.1.3] More data types ## Added diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 0000000000..dcd34d562a --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,41 @@ +# Contributor Covenant Code of Conduct + +## Our Pledge + +We as members, contributors, and leaders pledge to make participation in our community a harassment-free experience for everyone, regardless of age, body size, visible or invisible disability, ethnicity, sex characteristics, gender identity and expression, level of experience, education, socio-economic status, nationality, personal appearance, race, caste, color, religion, or sexual identity and orientation. + +We pledge to act and interact in ways that contribute to an open, welcoming, diverse, inclusive, and healthy community. + +## Our Standards + +Examples of behavior that contributes to a positive environment: + +* Demonstrating empathy and kindness toward other people +* Being respectful of differing opinions, viewpoints, and experiences +* Giving and gracefully accepting constructive feedback +* Accepting responsibility and apologizing to those affected by our mistakes +* Focusing on what is best not just for us as individuals, but for the overall community + +Examples of unacceptable behavior: + +* The use of sexualized language or imagery, and sexual attention or advances of any kind +* Trolling, insulting or derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information without their explicit permission +* Other conduct which could reasonably be considered inappropriate in a professional setting + +## Enforcement Responsibilities + +Project maintainers are responsible for clarifying and enforcing our standards of acceptable behavior and will take appropriate and fair corrective action in response to any behavior that they deem inappropriate, threatening, offensive, or harmful. + +## Scope + +This Code of Conduct applies within all community spaces, and also applies when an individual is officially representing the community in public spaces. + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behavior may be reported to the project maintainer at edward@weaviate.io. All complaints will be reviewed and investigated promptly and fairly. + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant](https://www.contributor-covenant.org), version 2.1. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5700fbb137..914883d820 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -8,7 +8,7 @@ Open source is at the heart of Verba. We appreciate feedback, ideas, and enhance ## 📚 Before You Begin -Before contributing, please take a moment to read through the [README](https://github.com/weaviate/Verba/README.md) and the [Technical Documentation](https://github.com/weaviate/Verba/TECHNICAL.md). These documents provide a comprehensive understanding of the project and are essential reading to ensure that we're all on the same page. Please note that the technical documentation is a work in progress and will be updated as we progress. +Before contributing, please take a moment to read through the [README](https://github.com/weaviate/Verba/README.md) and the [Backend Documentation](https://github.com/weaviate/Verba/goldenverba/README.md). These documents provide a comprehensive understanding of the project and are essential reading to ensure that we're all on the same page. ## 🐛 Reporting Issues diff --git a/LICENSE b/LICENSE index 90b8b9b57f..848f49a0ca 100644 --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright (c) 2020-2023, Weaviate B.V. +Copyright (c) 2020-2026, Weaviate B.V. All rights reserved. Redistribution and use in source and binary forms, with or without diff --git a/PYTHON_TUTORIAL.md b/PYTHON_TUTORIAL.md deleted file mode 100644 index 7c0461de42..0000000000 --- a/PYTHON_TUTORIAL.md +++ /dev/null @@ -1,70 +0,0 @@ -# Installing Python and Setting Up a Virtual Environment - -Before you can use Verba, you'll need to ensure that `Python >=3.10.0` is installed on your system and that you can create a virtual environment for a safer and cleaner project setup. - -## Installing Python - -Python is required to run Verba. If you don't have Python installed, follow these steps: - -### For Windows: - -Download the latest Python installer from the official Python website. -Run the installer and make sure to check the box that says `Add Python to PATH` during installation. - -### For macOS: - -You can install Python using Homebrew, a package manager for macOS, with the following command in the terminal: - -``` -/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)" -``` - -Then install Python: - -``` -brew install python -``` - -### For Linux: - -Python usually comes pre-installed on most Linux distributions. If it's not, you can install it using your distribution's package manager. You can read more about it [here](https://opensource.com/article/20/4/install-python-linux) - -## Setting Up a Virtual Environment - -It's recommended to use a virtual environment to avoid conflicts with other projects or system-wide Python packages. - -### Install the virtualenv package: - -First, ensure you have pip installed (it comes with Python if you're using version 3.4 and above). -Install virtualenv by running: - -``` -pip install virtualenv -``` - -### Create a Virtual Environment: - -Navigate to your project's directory in the terminal. -Run the following command to create a virtual environment named venv (you can name it anything you like): - -``` -python3 -m virtualenv venv -``` - -### Activate the Virtual Environment: - -- On Windows, activate the virtual environment by running: - -``` -venv\Scripts\activate.bat -``` - -- On macOS and Linux, activate it with: - -``` -source venv/bin/activate -``` - -Once your virtual environment is activated, you'll see its name in the terminal prompt. Now you're ready to install Verba using the steps provided in the Quickstart sections. - -> Remember to deactivate the virtual environment when you're done working with Verba by simply running deactivate in the terminal. diff --git a/README.md b/README.md index 9473c20ed9..6f3b8c9620 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ [![Weaviate](https://img.shields.io/static/v1?label=powered%20by&message=Weaviate%20%E2%9D%A4&color=green&style=flat-square)](https://weaviate.io/) [![PyPi downloads](https://static.pepy.tech/personalized-badge/goldenverba?period=total&units=international_system&left_color=grey&right_color=orange&left_text=pip%20downloads)](https://pypi.org/project/goldenverba/) [![Docker support](https://img.shields.io/badge/Docker_support-%E2%9C%93-4c1?style=flat-square&logo=docker&logoColor=white)](https://docs.docker.com/get-started/) [![Demo](https://img.shields.io/badge/Check%20out%20the%20demo!-yellow?&style=flat-square&logo=react&logoColor=white)](https://verba.weaviate.io/) -Welcome to Verba: The Golden RAGtriever, an community-driven open-source application designed to offer an end-to-end, streamlined, and user-friendly interface for Retrieval-Augmented Generation (RAG) out of the box. In just a few easy steps, explore your datasets and extract insights with ease, either locally with Ollama and Huggingface or through LLM providers such as Anthrophic, Cohere, and OpenAI. This project is built with and for the community, please be aware that it might not be maintained with the same urgency as other Weaviate production applications. Feel free to contribute to the project and help us make Verba even better! <3 +Welcome to Verba: The Golden RAGtriever, a community-driven open-source application designed to offer an end-to-end, streamlined, and user-friendly interface for Retrieval-Augmented Generation (RAG) out of the box. In just a few easy steps, explore your datasets and extract insights with ease, either locally with Ollama, HuggingFace, or LM Studio, or through LLM providers such as Anthropic, Cohere, OpenAI, DeepSeek, and more. This project is built with and for the community — feel free to contribute and help us make Verba even better! <3 ``` pip install goldenverba @@ -21,11 +21,13 @@ pip install goldenverba - [Weaviate](#weaviate) - [Ollama](#ollama) - [Unstructured](#unstructured) - - [AssemblyAI](#assemblyai) + - [Whisper (Audio/Video)](#whisper-audiovideo) - [OpenAI](#openai) - [HuggingFace](#huggingface) - [Groq](#groq) - - [Novita AI](#novitaai) + - [Novita AI](#novita) + - [DeepSeek](#deepseek) + - [LM Studio](#lm-studio) - [Quickstart: Deploy with pip](#how-to-deploy-with-pip) - [Quickstart: Build from Source](#how-to-build-from-source) - [Quickstart: Deploy with Docker](#how-to-install-verba-with-docker) @@ -48,37 +50,38 @@ Verba is a fully-customizable personal assistant utilizing [Retrieval Augmented ## Feature Lists -| 🤖 Model Support | Implemented | Description | -| --------------------------------- | ----------- | ------------------------------------------------------- | -| Ollama (e.g. Llama3) | ✅ | Local Embedding and Generation Models powered by Ollama | -| HuggingFace (e.g. MiniLMEmbedder) | ✅ | Local Embedding Models powered by HuggingFace | -| Cohere (e.g. Command R+) | ✅ | Embedding and Generation Models by Cohere | -| Anthrophic (e.g. Claude Sonnet) | ✅ | Embedding and Generation Models by Anthrophic | -| OpenAI (e.g. GPT4) | ✅ | Embedding and Generation Models by OpenAI | -| Groq (e.g. Llama3) | ✅ | Generation Models by Groq (LPU inference) | -| Novita AI (e.g. Llama3.3) | ✅ | Generation Models by Novita AI | -| Upstage (e.g. Solar) | ✅ | Embedding and Generation Models by Upstage | - -| 🤖 Embedding Support | Implemented | Description | -| -------------------- | ----------- | ---------------------------------------- | -| Weaviate | ✅ | Embedding Models powered by Weaviate | -| Ollama | ✅ | Local Embedding Models powered by Ollama | -| SentenceTransformers | ✅ | Embedding Models powered by HuggingFace | -| Cohere | ✅ | Embedding Models by Cohere | -| VoyageAI | ✅ | Embedding Models by VoyageAI | -| OpenAI | ✅ | Embedding Models by OpenAI | -| Upstage | ✅ | Embedding Models by Upstage | - -| 📁 Data Support | Implemented | Description | -| -------------------------------------------------------- | ----------- | ---------------------------------------------- | -| [UnstructuredIO](https://docs.unstructured.io/welcome) | ✅ | Import Data through Unstructured | -| [Firecrawl](https://www.firecrawl.dev/) | ✅ | Scrape and Crawl URL through Firecrawl | -| [UpstageDocumentParse](https://upstage.ai/) | ✅ | Parse Documents through Upstage Document AI | -| PDF Ingestion | ✅ | Import PDF into Verba | -| GitHub & GitLab | ✅ | Import Files from Github and GitLab | -| CSV/XLSX Ingestion | ✅ | Import Table Data into Verba | -| .DOCX | ✅ | Import .docx files | -| Multi-Modal (using [AssemblyAI](https://assemblyai.com)) | ✅ | Import and Transcribe Audio through AssemblyAI | +| 🤖 Model Support | Implemented | Description | +| --------------------------------- | ----------- | -------------------------------------------------------------------- | +| Ollama (e.g. Llama3) | ✅ | Local Embedding and Generation Models powered by Ollama | +| LM Studio | ✅ | Local Embedding and Generation Models via LM Studio (no API key) | +| HuggingFace (e.g. MiniLMEmbedder) | ✅ | Local Embedding Models powered by HuggingFace | +| Cohere (e.g. Command R+) | ✅ | Embedding and Generation Models by Cohere | +| Anthropic (e.g. Claude Sonnet) | ✅ | Generation Models by Anthropic | +| OpenAI (e.g. GPT4o) | ✅ | Embedding and Generation Models by OpenAI | +| DeepSeek (e.g. DeepSeek-R1) | ✅ | Generation Models by DeepSeek, including R1 reasoning with chain-of-thought display | +| Groq (e.g. Llama3) | ✅ | Generation Models by Groq (LPU inference) | +| Novita AI (e.g. Llama3.3) | ✅ | Generation Models by Novita AI | +| Upstage (e.g. Solar) | ✅ | Embedding and Generation Models by Upstage | + +| 🤖 Embedding Support | Implemented | Description | +| -------------------- | ----------- | -------------------------------------------------- | +| Weaviate | ✅ | Embedding Models powered by Weaviate | +| Ollama | ✅ | Local Embedding Models powered by Ollama | +| LM Studio | ✅ | Local Embedding Models via LM Studio (no API key) | +| SentenceTransformers | ✅ | Embedding Models powered by HuggingFace | +| Cohere | ✅ | Embedding Models by Cohere | +| VoyageAI | ✅ | Embedding Models by VoyageAI | +| OpenAI | ✅ | Embedding Models by OpenAI | +| Upstage | ✅ | Embedding Models by Upstage | + +| 📁 Data Support | Implemented | Description | +| -------------------------------------------------------- | ----------- | ------------------------------------------------------------ | +| [UnstructuredIO](https://docs.unstructured.io/welcome) | ✅ | Import Data through Unstructured (great for scanned PDFs) | +| PDF Ingestion | ✅ | Import PDF into Verba | +| GitHub & GitLab | ✅ | Import Files from Github and GitLab | +| CSV/XLSX Ingestion | ✅ | Import Table Data into Verba | +| .DOCX | ✅ | Import .docx files | +| Audio/Video (via [Whisper](https://github.com/guillaumekln/faster-whisper)) | ✅ | Transcribe audio and video locally — no API key required | | ✨ RAG Features | Implemented | Description | | ----------------------- | --------------- | ------------------------------------------------------------------------- | @@ -176,18 +179,22 @@ Below is a comprehensive list of the API keys and variables you may require: | COHERE_API_KEY | Your API Key | Get Access to [Cohere](https://cohere.com/) Models | | GROQ_API_KEY | Your Groq API Key | Get Access to [Groq](https://groq.com/) Models | | NOVITA_API_KEY | Your Novita API Key | Get Access to [Novita AI](https://novita.ai?utm_source=github_verba&utm_medium=github_readme&utm_campaign=github_link) Models | +| DEEPSEEK_API_KEY | Your DeepSeek API Key | Get Access to [DeepSeek](https://platform.deepseek.com/) Models | +| DEEPSEEK_BASE_URL | URL to DeepSeek instance (default: https://api.deepseek.com/v1) | Override the DeepSeek API endpoint | +| DEEPSEEK_MODEL | Model name (default: deepseek-chat) | Use `deepseek-chat` or `deepseek-reasoner` (R1) | +| LMSTUDIO_BASE_URL | URL to LM Studio (default: http://localhost:1234/v1) | Get Access to local models via [LM Studio](https://lmstudio.ai/) | +| LMSTUDIO_API_KEY | API key for LM Studio (optional) | Usually not required for local LM Studio | +| LMSTUDIO_MODEL | Generator model loaded in LM Studio | E.g. `lmstudio-community/Meta-Llama-3-8B-Instruct-GGUF` | +| LMSTUDIO_EMBEDDER_MODEL | Embedding model loaded in LM Studio | E.g. `nomic-ai/nomic-embed-text-v1.5-GGUF` | | OLLAMA_URL | URL to your Ollama instance (e.g. http://localhost:11434 ) | Get Access to [Ollama](https://ollama.com/) Models | | UNSTRUCTURED_API_KEY | Your API Key | Get Access to [Unstructured](https://docs.unstructured.io/welcome) Data Ingestion | | UNSTRUCTURED_API_URL | URL to Unstructured Instance | Get Access to [Unstructured](https://docs.unstructured.io/welcome) Data Ingestion | -| ASSEMBLYAI_API_KEY | Your API Key | Get Access to [AssemblyAI](https://assemblyai.com) Data Ingestion | | GITHUB_TOKEN | Your GitHub Token | Get Access to Data Ingestion via GitHub | | GITLAB_TOKEN | Your GitLab Token | Get Access to Data Ingestion via GitLab | -| FIRECRAWL_API_KEY | Your Firecrawl API Key | Get Access to Data Ingestion via Firecrawl | | VOYAGE_API_KEY | Your VoyageAI API Key | Get Access to Embedding Models via VoyageAI | | EMBEDDING_SERVICE_URL | URL to your Embedding Service Instance | Get Access to Embedding Models via [Weaviate Embedding Service](https://weaviate.io/developers/wcs/embeddings) | | EMBEDDING_SERVICE_KEY | Your Embedding Service Key | Get Access to Embedding Models via [Weaviate Embedding Service](https://weaviate.io/developers/wcs/embeddings) | -| UPSTAGE_API_KEY | Your Upstage API Key | Get Access to [Upstage](https://upstage.ai/) Models | -| UPSTAGE_BASE_URL | URL to Upstage instance | Models | +| UPSTAGE_API_KEY | Your Upstage API Key | Get Access to [Upstage](https://upstage.ai/) Embedding and Generation Models | | DEFAULT_DEPLOYMENT | Local, Weaviate, Custom, Docker | Set the default deployment mode | | SYSYEM_MESSAGE_PROMPT | Prompt text value | Default value starts with: "You are Verba, a chatbot for..." | | OLLAMA_MODEL | Your Ollama Model | Set the default Ollama model to use | @@ -239,9 +246,9 @@ Verba supports importing documents through Unstructured IO (e.g plain text, .pdf > UNSTRUCTURED_API_URL is set to `https://api.unstructuredapp.io/general/v0/general` by default -## AssemblyAI +## Whisper (Audio/Video) -Verba supports importing documents through AssemblyAI (audio files or audio from video files). To use them you need the `ASSEMBLYAI_API_KEY` environment variable. You can get it from [AssemblyAI](https://assemblyai.com) +Verba supports importing audio and video files (`.mp3`, `.wav`, `.m4a`, `.flac`, `.ogg`, `.mp4`, `.mov`, `.webm`, and more) via [faster-whisper](https://github.com/guillaumekln/faster-whisper), which runs locally with no API key or internet connection required. Install it with `pip install faster-whisper`. The model is downloaded automatically on first use (model files range from ~150 MB for `tiny` to ~3 GB for `large-v3`). ## OpenAI @@ -284,6 +291,29 @@ To use Groq LPUs as generation engine, you need to get an API key from [Groq](ht To use Novita AI as generation engine, you need to get an API key from [Novita AI](https://novita.ai/settings/key-management?utm_source=github_verba&utm_medium=github_readme&utm_campaign=github_link). +## DeepSeek + +Verba supports DeepSeek's models including the DeepSeek-R1 reasoning model. Get an API key from [DeepSeek Platform](https://platform.deepseek.com/). + +Set the `DEEPSEEK_API_KEY` environment variable, and optionally `DEEPSEEK_MODEL` to choose between `deepseek-chat` (default) and `deepseek-reasoner` (R1). When using `deepseek-reasoner`, the chain-of-thought reasoning process is displayed in a collapsible "Show Reasoning" section in the chat UI. + +``` +DEEPSEEK_API_KEY=your-deepseek-api-key +DEEPSEEK_MODEL=deepseek-reasoner # optional, defaults to deepseek-chat +``` + +## LM Studio + +Verba supports [LM Studio](https://lmstudio.ai/) for fully local inference with no API key required. Start LM Studio, load a model, and enable the local server (default: `http://localhost:1234/v1`). + +Both `LMStudioGenerator` and `LMStudioEmbedder` are available. Set environment variables to configure the endpoint and model names: + +``` +LMSTUDIO_BASE_URL=http://localhost:1234/v1 # optional, this is the default +LMSTUDIO_MODEL=your-loaded-chat-model +LMSTUDIO_EMBEDDER_MODEL=your-loaded-embedding-model +``` + # How to deploy with pip `Python >=3.10.0` @@ -433,7 +463,7 @@ Your contributions are always welcome! Feel free to contribute ideas, feedback, ### Project Architecture -You can learn more about Verba's architecture and implementation in its [technical documentation](./TECHNICAL.md) and [frontend documentation](./FRONTEND.md). It's recommended to have a look at them before making any contributions. +You can learn more about Verba's architecture and implementation in its [backend documentation](./goldenverba/README.md) and [frontend documentation](./FRONTEND.md). It's recommended to have a look at them before making any contributions. ## Known Issues diff --git a/TECHNICAL.md b/TECHNICAL.md deleted file mode 100644 index 509af5d098..0000000000 --- a/TECHNICAL.md +++ /dev/null @@ -1,55 +0,0 @@ -# Verba - Technical Documentation - -This technical documentation is intended for developers who want to understand the inner workings of Verba. Please note that this document might be uncomplete and missing some parts. If you encounter any issues or have questions, please feel free to open an issue. - -## FastAPI Server - -Verba is served through a FastAPI server. The server is serving the static frontend files through the specified port. If you're modifying the frontend, you will need to rebuild the static files again. The frontend is sending API calls to itself which the FastAPI server handles. The server can handle multiple client connections which are handled by the `ClientManager` class. - -### ClientManager - -`TODO` - -For handling large upload of files, the `BatchManager` class handles batches of data of a single file to merge it into a single file once all batches have been received. - -### BatchManager - -`TODO` - -### Websocket - -`TODO` - -## Automated Testing - -`TODO` - -## FAQ - -### How to control the position of context sent to the Generator to generate a response? - -Every `generator` class has a `prepare_messages` method. This method is used to format the messages that are sent to the LLM. The position of the context in the messages is important because it determines where the context is placed in the conversation. - -### How to upload a JSON file to Verba? - -## Verba JSON Structure - -A Verba Document can be created from a JSON object. The JSON object is converted to a Verba Document object and then uploaded to the vector database. Here's the general structure of a Verba Document (you can also find the implementation in the `Document.py` file): - -```python -{ - "title": "string", # The title of the document - "content": "string", # The content of the document - "extension": "string", # The extension of the document (Optional) - "fileSize": "number", # The size of the document in bytes (Optional) - "labels": "array", # The labels of the document (can be empty, used for filtering) - "source": "string", # The source of the document (can be an URL, optional) - "meta": "object", # The meta data of the document used internally - "metadata": "string" # Metadata information of the document, will be used in the embedding process -} -``` - -## Custom JSON Structure - -There is currently no support for custom JSON structure. Instead the whole JSON will simply be dumped into the content field of the Verba document. -There are plans to add support for custom JSON structure in the future. diff --git a/FRONTEND.md b/frontend/README.md similarity index 100% rename from FRONTEND.md rename to frontend/README.md diff --git a/frontend/package-lock.json b/frontend/package-lock.json index d9a0c32972..a9f27c7ef9 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -1,12 +1,12 @@ { "name": "verba", - "version": "2.1.0", + "version": "2.1.3", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "verba", - "version": "2.1.0", + "version": "2.1.3", "dependencies": { "@mdx-js/mdx": "^2.3.0", "@mdx-js/react": "^2.3.0", diff --git a/frontend/package.json b/frontend/package.json index a562414c62..479f43ff5c 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -1,6 +1,6 @@ { "name": "verba", - "version": "2.1.0", + "version": "2.1.3", "private": true, "scripts": { "dev": "next dev", diff --git a/goldenverba/.env.example b/goldenverba/.env.example index d0018bd12b..71f134bb5d 100644 --- a/goldenverba/.env.example +++ b/goldenverba/.env.example @@ -13,8 +13,6 @@ # EMBEDDING_SERVICE_URL= # EMBEDDING_SERVICE_KEY= -# FIRECRAWL_API_KEY= - # UNSTRUCTURED_API_KEY= # UNSTRUCTURED_API_URL=https://api.unstructuredapp.io/general/v0/general @@ -23,6 +21,9 @@ # OLLAMA_URL=http://localhost:11434 -# UPSTAGE_API_KEY= +# LMSTUDIO_BASE_URL=http://localhost:1234/v1 +# LMSTUDIO_API_KEY=lm-studio +# LMSTUDIO_MODEL=llama-3-8b-instruct +# LMSTUDIO_EMBEDDER_MODEL=text-embedding-nomic-embed-text-v1.5 # NOVITA_API_KEY= diff --git a/goldenverba/README.md b/goldenverba/README.md new file mode 100644 index 0000000000..3914993833 --- /dev/null +++ b/goldenverba/README.md @@ -0,0 +1,1051 @@ +# Verba Backend Developer Guide + +This guide is written for developers who want to contribute a new component to Verba — most commonly a new LLM provider (Generator), embedding provider (Embedder), file reader, or chunking strategy. It covers the plugin architecture, the data flow, the config system, and step-by-step walkthroughs for each component type. + +--- + +## Table of Contents + +1. [Architecture Overview](#architecture-overview) +2. [Project Structure](#project-structure) +3. [The Component Plugin System](#the-component-plugin-system) +4. [How to Add a New Generator](#how-to-add-a-new-generator) +5. [How to Add a New Embedder](#how-to-add-a-new-embedder) +6. [How to Add a New Reader](#how-to-add-a-new-reader) +7. [How to Add a New Chunker](#how-to-add-a-new-chunker) +8. [The Config System](#the-config-system) +9. [WebSocket Protocol](#websocket-protocol) +10. [Verba Document JSON Format](#verba-document-json-format) +11. [Running Locally for Development](#running-locally-for-development) +12. [API Reference](#api-reference) + +--- + +## Architecture Overview + +### Component Plugin System + +Verba uses a simple plugin architecture. Five component categories exist; each category has a manager class that holds a dict of registered instances. The frontend reads available components via REST, the user picks one, and its config is serialised into every subsequent request. + +``` ++----------------------------------------------------------+ +| Frontend (Next.js) | +| - Picks Reader / Chunker / Embedder / Retriever / | +| Generator from dropdowns populated by /api/health | +| and /api/get_rag_config | ++---------------------------+------------------------------+ + | REST + WebSocket ++---------------------------v------------------------------+ +| FastAPI (goldenverba/server/api.py) | +| | +| VerbaManager | +| +----------------+ +----------------+ | +| | ReaderManager | | ChunkerManager | | +| +----------------+ +----------------+ | +| +------------------+ +------------------+ | +| | EmbeddingManager | | RetrieverManager | | +| +------------------+ +------------------+ | +| +------------------+ | +| | GeneratorManager | | +| +------------------+ | +| | +| WeaviateManager (all Weaviate I/O lives here) | ++---------------------------+------------------------------+ + | ++---------------------------v------------------------------+ +| Weaviate | +| Collections: VERBA_DOCUMENTS, VERBA_CONFIGURATION, | +| VERBA_SUGGESTIONS | ++----------------------------------------------------------+ +``` + +### The Five Component Types + +| Type | Base class | Purpose | +|---|---|---| +| Reader | `Reader` | Turn a file/URL upload into `Document` objects | +| Chunker | `Chunker` | Split `Document` objects into `Chunk` objects | +| Embedder | `Embedding` | Vectorise each chunk's text into `list[float]` | +| Retriever | `Retriever` | Query Weaviate and build the context string | +| Generator | `Generator` | Stream an LLM response token by token | + +### Full Data Flow: File Upload to Streamed Answer + +``` +Browser uploads file + | + v +[WebSocket /ws/import_files] + | + BatchManager reassembles chunks into FileConfig + | + v +VerbaManager.import_document(client, fileConfig) + | + +--- ReaderManager.load() --> list[Document] + | (calls Reader.load()) + | + +--- ChunkerManager.chunk() --> list[Document] with .chunks filled + | (calls Chunker.chunk()) + | + +--- EmbeddingManager.vectorize() --> list[Document] with chunk.vector filled + | (calls Embedding.vectorize() in batches of max_batch_size) + | also computes PCA(3) for 3-D visualisation + | + +--- WeaviateManager.import_document() + stores Document + Chunks in VERBA_DOCUMENTS collection + +------- query time ------- + +Browser sends query text + | + v +[POST /api/query] + | + VerbaManager.retrieve_chunks() + | + +--- EmbeddingManager.vectorize_query() --> query vector + | + +--- RetrieverManager.retrieve() --> (documents, context_str) + (Retriever hits Weaviate with vector + filters) + | + v +[WebSocket /ws/generate_stream] + | + GeneratorManager.generate_stream() + | + Generator.generate_stream() -- async generator + | + yields {"message": "", "finish_reason": None | "stop"} + | + WebSocket sends each dict to the browser in real time +``` + +### Production vs Local Mode + +The `VERBA_PRODUCTION` environment variable controls which components are available at startup. The filter lives in each manager file (`embedding_manager.py` and `generator_manager.py`). + +| Value | Effect | +|---|---| +| unset / anything else | All components including Ollama, SentenceTransformers, Groq, Novita AI | +| `"Production"` | Local-only components excluded (Ollama, SentenceTransformers, Groq, Novita AI) | +| `"Demo"` | Same component set as Production but write endpoints return early | + +The `production` string is exposed to the frontend via `GET /api/health` so the UI can hide destructive actions. + +--- + +## Project Structure + +``` +Verba/ +├── goldenverba/ # Entire Python package +│ ├── __init__.py +│ ├── components/ +│ │ ├── verba_manager.py # VerbaManager — orchestrates the full pipeline +│ │ ├── client_manager.py # ClientManager — Weaviate connection pooling +│ │ ├── interfaces.py # Abstract base classes (VerbaComponent, Reader, +│ │ │ # Chunker, Embedding, Retriever, Generator) +│ │ ├── weaviate_manager.py # All direct Weaviate I/O (collections, CRUD, search) +│ │ ├── types.py # InputConfig Pydantic model +│ │ ├── util.py # get_environment(), get_token() helpers +│ │ ├── document.py # Document dataclass +│ │ ├── chunk.py # Chunk dataclass +│ │ ├── reader/ +│ │ │ ├── reader_manager.py # readers list + ReaderManager +│ │ │ ├── BasicReader.py # Default: txt, pdf, docx, csv, xlsx, … +│ │ │ ├── GitReader.py # Clones a Git repo and reads files +│ │ │ ├── HTMLReader.py # Fetches a URL and strips HTML +│ │ │ ├── UnstructuredAPI.py # Unstructured.io API reader +│ │ │ └── WhisperReader.py # Local audio/video transcription via faster-whisper +│ │ ├── chunking/ +│ │ │ ├── chunker_manager.py # chunkers list + ChunkerManager +│ │ │ ├── TokenChunker.py # Token-count window with overlap +│ │ │ ├── SentenceChunker.py +│ │ │ ├── RecursiveChunker.py +│ │ │ ├── SemanticChunker.py # Embedding-aware chunker +│ │ │ ├── MarkdownChunker.py +│ │ │ ├── HTMLChunker.py +│ │ │ ├── CodeChunker.py +│ │ │ └── JSONChunker.py +│ │ ├── embedding/ +│ │ │ ├── embedding_manager.py # embedders list (+ production filter) + EmbeddingManager +│ │ │ ├── OpenAIEmbedder.py +│ │ │ ├── CohereEmbedder.py +│ │ │ ├── OllamaEmbedder.py +│ │ │ ├── WeaviateEmbedder.py +│ │ │ ├── VoyageAIEmbedder.py +│ │ │ ├── SentenceTransformersEmbedder.py +│ │ │ ├── UpstageEmbedder.py +│ │ │ └── LMStudioEmbedder.py +│ │ ├── generation/ +│ │ │ ├── generator_manager.py # generators list (+ production filter) + GeneratorManager +│ │ │ ├── OpenAIGenerator.py +│ │ │ ├── AnthrophicGenerator.py +│ │ │ ├── CohereGenerator.py +│ │ │ ├── OllamaGenerator.py +│ │ │ ├── GroqGenerator.py +│ │ │ ├── NovitaGenerator.py +│ │ │ ├── DeepSeekGenerator.py +│ │ │ ├── UpstageGenerator.py +│ │ │ └── LMStudioGenerator.py +│ │ └── retriever/ +│ │ ├── retriever_manager.py # retrievers list + RetrieverManager +│ │ └── WindowRetriever.py # Default retriever (window expansion) +│ ├── server/ +│ │ ├── api.py # FastAPI app, all endpoints, WebSocket handlers +│ │ ├── types.py # Pydantic request/response models +│ │ ├── helpers.py # LoggerManager, BatchManager +│ │ ├── cli.py # `verba start` / `verba reset` Click commands +│ │ └── frontend/out/ # Pre-built Next.js static files (served by FastAPI) +│ └── tests/ +│ ├── components/ +│ │ ├── test_chunk.py # Chunk model tests +│ │ ├── test_document.py # Document model tests +│ │ ├── chunking/ +│ │ │ └── test_chunkers.py # TokenChunker, SentenceChunker, MarkdownChunker +│ │ └── weaviate/ +│ │ └── test_weaviate_manager.py # Integration tests (opt-in) +│ └── server/ +│ ├── test_api.py # FastAPI endpoint tests +│ └── test_helpers.py # BatchManager, LoggerManager tests +├── frontend/ # Next.js source (separate dev server in development) +├── setup.py # Package metadata + dependencies +├── requirements.txt # Mirror of setup.py install_requires +└── .env # Local secrets (never committed) +``` + +--- + +## The Component Plugin System + +### Base Classes (`goldenverba/components/interfaces.py`) + +Every component inherits from `VerbaComponent`: + +```python +class VerbaComponent: + def __init__(self): + self.name = "" # Unique display name (used as dict key) + self.requires_env = [] # Env vars that must be set for availability + self.requires_library = [] # Python packages that must be importable + self.description = "" # Shown in the UI + self.config = {} # Dict[str, InputConfig] — UI-configurable fields + self.type = "" # Internal type tag +``` + +`get_meta()` serialises the component into the dict that the frontend consumes. It calls `check_available()` which checks `requires_env` against current environment variables and `requires_library` against installed packages. + +### InputConfig — Configurable UI Fields (`goldenverba/components/types.py`) + +```python +class InputConfig(BaseModel): + type: Literal["number", "text", "dropdown", "password", "bool", "multi", "textarea"] + value: Union[int, str, bool] # Current (or default) value + description: str # Help text shown below the field + values: list[str] # Options — only used when type == "dropdown" or "multi" +``` + +Each key in `self.config` becomes a labelled input in the UI. The frontend sends the complete config dict back with every request so components always receive whatever the user last set. + +| `type` value | Rendered as | When to use | +|---|---|---| +| `"text"` | Single-line text box | Base URL, model name overrides | +| `"password"` | Masked input | API keys that are not in env | +| `"number"` | Numeric spinner | Token counts, temperature | +| `"dropdown"` | Select menu | Model selection (populate `values` list) | +| `"bool"` | Toggle | Feature flags | +| `"textarea"` | Multi-line text | System prompts | +| `"multi"` | Multi-select | Tag-style multi-value fields | + +### How `get_meta()` Works + +`VerbaManager` calls `get_meta(envs, libs)` on every registered component at connect time to build the `RAGConfig` that is returned to the frontend. `envs` is a snapshot of `os.environ` and `libs` is a dict of `{package_name: bool}` built by `VerbaManager.verify_installed_libraries()`. + +A component with `requires_env = ["OPENAI_API_KEY"]` will have `"available": false` in the UI unless that variable is set, preventing the user from selecting an unusable component. + +### How Components Are Registered + +Each component category owns its list in its manager file. To add a new component, edit only that one file: + +| Category | Manager file | +|---|---| +| Readers | `goldenverba/components/reader/reader_manager.py` | +| Chunkers | `goldenverba/components/chunking/chunker_manager.py` | +| Embedders | `goldenverba/components/embedding/embedding_manager.py` | +| Retrievers | `goldenverba/components/retriever/retriever_manager.py` | +| Generators | `goldenverba/components/generation/generator_manager.py` | + +Example — `generator_manager.py` (simplified): + +```python +# generator_manager.py +from goldenverba.components.generation.OpenAIGenerator import OpenAIGenerator +from goldenverba.components.generation.OllamaGenerator import OllamaGenerator +# ... + +_is_production = os.getenv("VERBA_PRODUCTION") == "Production" + +generators = [ + g for g in [ + OllamaGenerator(), + OpenAIGenerator(), + # ... + ] + if not _is_production or g.name not in {"Ollama", "Groq", "Novita AI"} +] + +class GeneratorManager: + def __init__(self): + self.generators = {g.name: g for g in generators} +``` + +--- + +## How to Add a New Generator + +This walkthrough adds a fictional "Mistral" generator backed by the Mistral AI chat completions API. + +### Step 1 — Create the file + +``` +goldenverba/components/generation/MistralGenerator.py +``` + +### Step 2 — Write the class + +```python +import os +import json +import httpx +from goldenverba.components.interfaces import Generator +from goldenverba.components.types import InputConfig +from goldenverba.components.util import get_environment, get_token + +class MistralGenerator(Generator): + def __init__(self): + super().__init__() + self.name = "Mistral" # Must be unique + self.description = "Mistral AI chat completions" + self.context_window = 32000 + + # Offer a model dropdown. Fallback to a static list if no key yet. + api_key = get_token("MISTRAL_API_KEY") + models = self._fetch_models(api_key) if api_key else ["mistral-small", "mistral-large-latest"] + default_model = os.getenv("MISTRAL_MODEL", models[0]) + + self.config["Model"] = InputConfig( + type="dropdown", + value=default_model, + description="Select a Mistral model", + values=models, + ) + + # Only show the API key field if the env var is not already set. + if api_key is None: + self.config["API Key"] = InputConfig( + type="password", + value="", + description="Mistral API key — or set MISTRAL_API_KEY env var", + values=[], + ) + + async def generate_stream(self, config: dict, query: str, context: str, conversation: list[dict] = []): + model = config["Model"].value + api_key = get_environment(config, "API Key", "MISTRAL_API_KEY", "No Mistral API key found") + system_message = config["System Message"].value # inherited from Generator base class + + messages = self.prepare_messages(query, context, conversation, system_message) + + headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } + body = {"model": model, "messages": messages, "stream": True} + + async with httpx.AsyncClient() as client: + async with client.stream( + "POST", + "https://api.mistral.ai/v1/chat/completions", + json=body, + headers=headers, + timeout=None, + ) as response: + async for line in response.aiter_lines(): + if line.startswith("data: "): + if line.strip() == "data: [DONE]": + break + data = json.loads(line[6:]) + choice = data["choices"][0] + delta = choice.get("delta", {}) + if "content" in delta: + yield {"message": delta["content"], "finish_reason": choice.get("finish_reason")} + elif choice.get("finish_reason"): + yield {"message": "", "finish_reason": choice["finish_reason"]} + + def prepare_messages(self, query: str, context: str, conversation: list[dict], system_message: str) -> list[dict]: + messages = [{"role": "system", "content": system_message}] + for msg in conversation: + messages.append({"role": msg.type, "content": msg.content}) + messages.append({"role": "user", "content": f"Answer this query: '{query}' with this context: {context}"}) + return messages + + def _fetch_models(self, api_key: str) -> list[str]: + try: + import requests + r = requests.get("https://api.mistral.ai/v1/models", headers={"Authorization": f"Bearer {api_key}"}) + r.raise_for_status() + return [m["id"] for m in r.json()["data"]] + except Exception: + return ["mistral-small", "mistral-large-latest"] +``` + +### Step 3 — Register in `generator_manager.py` + +Open `goldenverba/components/generation/generator_manager.py` and add two lines: + +```python +# At the top with the other generator imports: +from goldenverba.components.generation.MistralGenerator import MistralGenerator + +# In the generators list: +generators = [ + g for g in [ + OllamaGenerator(), + OpenAIGenerator(), + MistralGenerator(), # <-- add here + ... + ] + if not _is_production or g.name not in {"Ollama", "Groq", "Novita AI"} +] +``` + +If the generator requires local infrastructure (no cloud API), also add its name to the production exclusion set. + +### Step 4 — Add environment variable documentation + +Document the new variable in your PR description and, if relevant, in `.env.example`: + +``` +MISTRAL_API_KEY=your-key-here +MISTRAL_MODEL=mistral-large-latest # optional +``` + +### Key rules for `generate_stream` + +- It must be an `async` generator (use `yield`, not `return`). +- Every yielded dict must have exactly two keys: `"message"` (str) and `"finish_reason"` (str or `None`). +- Signal end-of-stream by yielding `{"message": "", "finish_reason": "stop"}`. +- Never raise inside the generator after streaming has begun — send a terminal chunk instead. +- `config` is a `dict[str, InputConfig]`. Access values with `config["Key"].value`. + +### `get_environment` vs `get_token` + +```python +from goldenverba.components.util import get_environment, get_token + +# get_token: reads ONLY from os.environ, returns None if missing/empty +api_key = get_token("MISTRAL_API_KEY") # use in __init__ to check availability + +# get_environment: checks config dict first, then falls back to env var, raises if neither found +api_key = get_environment(config, "API Key", "MISTRAL_API_KEY", "No Mistral API key") +# use in generate_stream / vectorize where config is available +``` + +The pattern — check with `get_token` in `__init__` to decide whether to show the password field, then resolve with `get_environment` at call time — means users can either set an env var before starting Verba or paste the key directly into the UI. + +--- + +## How to Add a New Embedder + +This walkthrough adds a fictional "VoyageMini" embedder. + +### Step 1 — Create the file + +``` +goldenverba/components/embedding/VoyageMiniEmbedder.py +``` + +### Step 2 — Write the class + +```python +import aiohttp +from goldenverba.components.interfaces import Embedding +from goldenverba.components.types import InputConfig +from goldenverba.components.util import get_environment, get_token + +class VoyageMiniEmbedder(Embedding): + def __init__(self): + super().__init__() + self.name = "VoyageMini" + self.description = "Voyage AI mini embedding model" + self.max_batch_size = 64 # Voyage mini accepts up to 128, set conservatively + + self.config["Model"] = InputConfig( + type="dropdown", + value="voyage-3-lite", + description="Voyage embedding model", + values=["voyage-3-lite", "voyage-3"], + ) + + if get_token("VOYAGE_API_KEY") is None: + self.config["API Key"] = InputConfig( + type="password", + value="", + description="Voyage AI API key — or set VOYAGE_API_KEY env var", + values=[], + ) + + async def vectorize(self, config: dict, content: list[str]) -> list[list[float]]: + """ + Receives a batch of strings (already split by EmbeddingManager based on + self.max_batch_size) and returns one embedding vector per string. + """ + model = config["Model"].value + api_key = get_environment(config, "API Key", "VOYAGE_API_KEY", "No Voyage API key found") + + payload = {"input": content, "model": model} + headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} + + async with aiohttp.ClientSession() as session: + async with session.post( + "https://api.voyageai.com/v1/embeddings", + json=payload, + headers=headers, + timeout=aiohttp.ClientTimeout(total=60), + ) as resp: + resp.raise_for_status() + data = await resp.json() + return [item["embedding"] for item in data["data"]] +``` + +### Step 3 — Register in `embedding_manager.py` + +```python +# goldenverba/components/embedding/embedding_manager.py + +from goldenverba.components.embedding.VoyageMiniEmbedder import VoyageMiniEmbedder + +embedders = [ + e for e in [ + ..., + VoyageMiniEmbedder(), # <-- add here + ] + if not _is_production or e.name not in {"Ollama", "SentenceTransformers"} +] +``` + +If the embedder requires local infrastructure, add its name to the production exclusion set. + +### Embedder contract + +- `vectorize(config, content)` receives a **batch** of strings (at most `self.max_batch_size` items). +- It must return `list[list[float]]` — one inner list per input string, all the same length. +- `EmbeddingManager` calls `vectorize` in parallel for each batch via `asyncio.gather`, then flattens and verifies counts. +- Set `self.max_batch_size` to the largest batch the upstream API accepts reliably. + +--- + +## Available Readers + +| Reader | Type | Supported Formats | Env Var | pip Package | +|---|---|---|---|---| +| **Default** (`BasicReader`) | FILE | `.txt .md .mdx .json .html .css .py .ts .tsx .js .go .rs .swift .kt .java .c .cpp .h .hpp .vue .svelte .astro .php .rb` + `.pdf .docx .pptx .csv .xlsx .xls` | — | `pypdf`, `python-docx`, `pandas`, `openpyxl` | +| **HTML** (`HTMLReader`) | URL | Any HTTP/HTTPS URL — optional Markdown conversion, recursive same-domain crawl | — | `beautifulsoup4`, `markdownify` | +| **Git** (`GitReader`) | URL | All formats supported by Default reader, fetched from GitHub or GitLab via API | `GITHUB_TOKEN` / `GITLAB_TOKEN` | — | +| **Unstructured IO** (`UnstructuredReader`) | FILE | 40+ formats via Unstructured.io API (PDF, DOCX, images, …); strategies: `auto`, `hi_res`, `ocr_only`, `fast` | `UNSTRUCTURED_API_KEY` | — | +| **Whisper** (`WhisperReader`) | FILE | 35+ audio/video formats (`.mp3 .wav .m4a .flac .ogg .aac .mp4 .mov .webm` …) — transcribed locally, no API key required | — | `faster-whisper` | + +**Reader type** tells the UI whether to show a file picker (`FILE`) or a URL input (`URL`). + +--- + +## How to Add a New Reader + +### Step 1 — Create the file + +``` +goldenverba/components/reader/MyReader.py +``` + +### Step 2 — Write the class + +```python +import base64 +from goldenverba.components.interfaces import Reader +from goldenverba.components.document import Document, create_document +from goldenverba.server.types import FileConfig + +class MyReader(Reader): + def __init__(self): + super().__init__() + self.name = "MyReader" + self.description = "Reads .xyz files" + self.type = "FILE" # "FILE" or "URL" + self.extension = [".xyz"] # Accepted extensions shown in UI + + async def load(self, config: dict, fileConfig: FileConfig) -> list[Document]: + """ + fileConfig.content is base64-encoded file bytes when fileConfig.isURL is False. + Must return a list of Document objects. + """ + raw = base64.b64decode(fileConfig.content) + text = raw.decode("utf-8") + + # create_document is a convenience helper that populates all required fields + return [create_document(text, fileConfig)] +``` + +### `FileConfig` fields relevant to readers + +| Field | Type | Description | +|---|---|---| +| `filename` | `str` | Original filename | +| `extension` | `str` | File extension without leading dot | +| `content` | `str` | Base64-encoded file bytes (or raw URL string when `isURL=True`) | +| `isURL` | `bool` | True when the source is a URL rather than an uploaded file | +| `labels` | `list[str]` | User-supplied labels to attach to the document | +| `metadata` | `str` | Free-form string for extra metadata | +| `overwrite` | `bool` | Whether to replace an existing document with the same name | +| `rag_config` | `dict` | Full RAG config including the reader's own config fields | + +### Returning multiple documents + +Readers may return more than one `Document` from a single file (e.g. a zip archive or a URL that fans out to multiple pages). When more than one document is returned, `VerbaManager` processes each concurrently. + +### Step 3 — Register in `reader_manager.py` + +```python +# goldenverba/components/reader/reader_manager.py + +from goldenverba.components.reader.MyReader import MyReader + +readers = [ + BasicReader(), + MyReader(), # <-- add here + ... +] +``` + +--- + +## How to Add a New Chunker + +### Step 1 — Create the file + +``` +goldenverba/components/chunking/MyChunker.py +``` + +### Step 2 — Write the class + +```python +from goldenverba.components.interfaces import Chunker, Embedding +from goldenverba.components.chunk import Chunk +from goldenverba.components.document import Document +from goldenverba.components.types import InputConfig + +class MyChunker(Chunker): + def __init__(self): + super().__init__() + self.name = "MyChunker" + self.description = "Splits on paragraph boundaries" + self.config["Max Paragraphs"] = InputConfig( + type="number", + value=3, + description="Maximum paragraphs per chunk", + values=[], + ) + + async def chunk( + self, + config: dict, + documents: list[Document], + embedder: Embedding | None = None, + embedder_config: dict | None = None, + ) -> list[Document]: + max_p = int(config["Max Paragraphs"].value) + + for doc in documents: + if doc.chunks: # skip already-chunked documents + continue + paragraphs = doc.content.split("\n\n") + groups = [paragraphs[i:i+max_p] for i in range(0, len(paragraphs), max_p)] + for idx, group in enumerate(groups): + text = "\n\n".join(group) + doc.chunks.append(Chunk( + content=text, + chunk_id=idx, + start_i=0, + end_i=len(text), + content_without_overlap=text, + )) + return documents +``` + +### Chunker contract + +- Receives a list of `Document` objects that already have `doc.content` and `doc.spacy_doc` populated. +- Must populate `doc.chunks` with `Chunk` instances. +- Skip documents that already have chunks (`if doc.chunks: continue`). +- `embedder` and `embedder_config` are only provided when the chunker needs to call the embedder itself (e.g. `SemanticChunker` uses them to split on embedding similarity boundaries). Most chunkers ignore them. +- Returns the same list of documents. + +### Step 3 — Register in `chunker_manager.py` + +```python +# goldenverba/components/chunking/chunker_manager.py + +from goldenverba.components.chunking.MyChunker import MyChunker + +chunkers = [ + TokenChunker(), + MyChunker(), # <-- add here + ... +] +``` + +--- + +## The Config System + +### Where Config Lives + +RAG configuration is persisted in the `VERBA_CONFIGURATION` Weaviate collection as a JSON blob under a fixed UUID (`VerbaManager.rag_config_uuid`). When a client connects, `VerbaManager.load_rag_config()` reads this blob and merges it with the current component definitions to produce the full `RAGConfig` sent to the frontend. + +### Config Shape + +The frontend always works with a `RAGConfig` object: + +``` +RAGConfig + Reader: RAGComponentClass { selected: "Default", components: { "Default": RAGComponentConfig, ... } } + Chunker: RAGComponentClass { selected: "Token", components: { ... } } + Embedder: RAGComponentClass { selected: "OpenAI", components: { ... } } + Retriever: RAGComponentClass { selected: "Window", components: { ... } } + Generator: RAGComponentClass { selected: "OpenAI", components: { ... } } +``` + +Each `RAGComponentConfig` carries the component's `config` dict (key → `ConfigSetting`). The `ConfigSetting` type mirrors `InputConfig` exactly — when the user edits a field and saves, the new `value` is written back into this structure and persisted via `POST /api/set_rag_config`. + +### How Config Is Passed to Components + +Every manager method extracts the config before calling the component: + +```python +# ChunkerManager.chunk() — typical pattern +config = fileConfig.rag_config["Chunker"].components[chunker].config +await self.chunkers[chunker].chunk(config=config, documents=documents, ...) +``` + +Inside a component, read values with `config["Key"].value`. The value is whatever the user last set in the UI, or the default you put in `__init__`. + +### Config Fields Added Conditionally + +The standard pattern is to omit a config field if the corresponding env var is already set. This keeps the UI clean for users who configure via env vars: + +```python +if get_token("MY_API_KEY") is None: + self.config["API Key"] = InputConfig(type="password", ...) +``` + +If the env var is set, the field never appears and `get_environment(config, "API Key", "MY_API_KEY", "...")` will fall through to reading the env var directly. + +--- + +## WebSocket Protocol + +### `/ws/generate_stream` — Streaming Generation + +**Client sends** (JSON text frame, validated against `GeneratePayload`): +```json +{ + "query": "What is RAG?", + "context": "RAG stands for...", + "conversation": [ + {"type": "user", "content": "Hello"}, + {"type": "assistant", "content": "Hi there!"} + ], + "rag_config": { "" } +} +``` + +**Server yields** (one JSON text frame per token): +```json +{"message": "Retrieval", "finish_reason": null} +{"message": "-", "finish_reason": null} +{"message": "Augmented", "finish_reason": null} +{"message": " Generation", "finish_reason": null} +{"message": "", "finish_reason": "stop", "full_text": "Retrieval-Augmented Generation"} +``` + +The final frame where `finish_reason == "stop"` also carries `full_text` (the entire concatenated response). On error the server sends `{"message": "", "finish_reason": "stop", "full_text": ""}`. + +The connection stays open after each exchange so the user can send a follow-up query without reconnecting. + +### `/ws/import_files` — Batch File Upload + +Large `FileConfig` objects are split into chunks by the frontend before sending. The `BatchManager` on the server side reassembles them. + +**Client sends** (one frame per batch, validated against `DataBatchPayload`): +```json +{ + "chunk": "", + "isLastChunk": false, + "total": 3, + "fileID": "abc-123", + "order": 0, + "credentials": {"deployment": "Docker", "url": "", "key": ""} +} +``` + +Fields: +- `total` — total number of chunks for this file transfer. +- `order` — 0-indexed position of this chunk. +- `isLastChunk` — when true the BatchManager cleans up regardless of completeness (handles partial failures). +- `fileID` — unique identifier for this file upload session. + +**Server sends** (status frames via `LoggerManager`): +```json +{"fileID": "abc-123", "status": "LOADING", "message": "Loaded document.pdf", "took": 0.12} +{"fileID": "abc-123", "status": "CHUNKING", "message": "Split into 47 chunks", "took": 0.44} +{"fileID": "abc-123", "status": "EMBEDDING", "message": "Vectorized all chunks", "took": 3.21} +{"fileID": "abc-123", "status": "INGESTING", "message": "Imported into Weaviate", "took": 0.8} +{"fileID": "abc-123", "status": "DONE", "message": "Import completed successfully", "took": 4.6} +``` + +On failure: `{"fileID": "...", "status": "ERROR", "message": "", "took": 0}`. + +Full status lifecycle: `READY → STARTING → LOADING → CHUNKING → EMBEDDING → INGESTING → DONE` (or `ERROR`). + +When a Reader returns multiple documents (e.g. a URL reader that fans out), the server sends a `CreateNewDocument` frame so the UI can track each sub-document separately: +```json +{"new_file_id": "abc-123", "filename": "page-title", "original_file_id": "abc-123"} +``` + +--- + +## Verba Document JSON Format + +A Verba Document can be created directly from a JSON file upload. The JSON is converted to a `Document` object and imported into Weaviate. Structure: + +```json +{ + "title": "string", // Display title of the document + "content": "string", // Main text content + "extension": "string", // File extension, e.g. ".txt" (optional) + "fileSize": 0, // Size in bytes (optional) + "labels": [], // Array of strings used for filtering (can be empty) + "source": "string", // Source URL or path (optional) + "meta": {}, // Internal metadata object — populated by the pipeline + "metadata": "string" // Free-form string prepended to chunks during embedding +} +``` + +The `metadata` field is prepended to each chunk's content before vectorisation, making it useful for injecting document-level context (author, date, category) into every chunk's embedding without cluttering the displayed content. + +`meta` is managed by the pipeline (Reader, Chunker, Embedder each write their config into it) and should not be set manually in uploaded JSON. + +--- + +## Running Locally for Development + +### Prerequisites + +- Python 3.10, 3.11, or 3.12 (3.13+ not yet supported) +- Node.js 18+ (only needed if modifying the frontend) +- A running Weaviate instance — Docker Compose (recommended) or Weaviate Cloud + +### Python Environment + +```bash +git clone https://github.com/weaviate/Verba.git +cd Verba +python -m venv .venv +source .venv/bin/activate # Windows: .venv\Scripts\activate + +pip install -e ".[dev]" + +# Install the spaCy model required by BasicReader / TokenChunker +python -m spacy download en_core_web_sm +``` + +### Environment Variables + +Create a `.env` file in the project root: + +```bash +# Weaviate connection +WEAVIATE_URL_VERBA=https://your-cluster.weaviate.network # for Weaviate Cloud +WEAVIATE_API_KEY_VERBA=your-weaviate-key + +# API keys for providers you want to test +OPENAI_API_KEY=sk-... +ANTHROPIC_API_KEY=sk-ant-... + +# Optional: pre-select models +OPENAI_MODEL=gpt-4o +OPENAI_EMBED_MODEL=text-embedding-3-small + +# Optional: override system prompt +SYSYEM_MESSAGE_PROMPT="You are a helpful assistant." + +# Optional: run in production component set (excludes Ollama, SentenceTransformers, Groq, Novita AI) +# VERBA_PRODUCTION=Production +``` + +### Running the Backend + +```bash +verba start --port 8000 --host localhost +# or: +python -m goldenverba.server.cli start --port 8000 +``` + +Access the bundled frontend at `http://localhost:8000`. + +### Running the Frontend Separately (for UI development) + +```bash +# Terminal 1 — backend +verba start --port 8000 --prod + +# Terminal 2 — frontend dev server +cd frontend +npm install +npm run dev # starts Next.js on http://localhost:3000 +``` + +The Next.js dev server proxies API calls to port 8000 (configured in `frontend/next.config.js`). + +### Resetting Weaviate Collections + +```bash +# Reset only config (keeps documents) +verba reset --deployment Docker + +# Full wipe of all VERBA_ collections +verba reset --deployment Docker --full_reset True + +# Against a Weaviate Cloud cluster +verba reset --url https://my-cluster.weaviate.network --api_key <key> --deployment Weaviate +``` + +### Running Tests + +```bash +# Unit tests (no Weaviate required) +pytest goldenverba/tests -v + +# Integration tests (requires a live Weaviate instance) +WEAVIATE_INTEGRATION=1 pytest goldenverba/tests/components/weaviate -v +# or against Weaviate Cloud: +WEAVIATE_TEST_URL=https://my-cluster.weaviate.network WEAVIATE_TEST_KEY=<key> pytest goldenverba/tests/components/weaviate -v +``` + +### Code Style + +Format with Black before submitting a PR: + +```bash +black goldenverba/ +``` + +--- + +## API Reference + +All endpoints accept and return JSON. Every endpoint except `/api/health` requires an `Origin` header matching the server's base URL (same-origin middleware). Credentials are embedded in request bodies rather than headers. + +### Health and Connection + +| Method | Path | Request body | Response | +|---|---|---|---| +| `GET` | `/api/health` | — | `{message, production, gtag, deployments, default_deployment}` | +| `POST` | `/api/connect` | `ConnectPayload {credentials, port}` | `{connected, error, rag_config, user_config, theme, themes}` | + +### Configuration + +| Method | Path | Request body | Response | +|---|---|---|---| +| `POST` | `/api/get_rag_config` | `Credentials` | `{rag_config, error}` | +| `POST` | `/api/set_rag_config` | `SetRAGConfigPayload {rag_config, credentials}` | `{status}` | +| `POST` | `/api/get_user_config` | `Credentials` | `{user_config, error}` | +| `POST` | `/api/set_user_config` | `SetUserConfigPayload {user_config, credentials}` | `{status, status_msg}` | +| `POST` | `/api/get_theme_config` | `Credentials` | `{theme, themes, error}` | +| `POST` | `/api/set_theme_config` | `SetThemeConfigPayload {theme, themes, credentials}` | `{status}` | + +### RAG / Query + +| Method | Path | Request body | Response | +|---|---|---|---| +| `POST` | `/api/query` | `QueryPayload {query, RAG, labels, documentFilter, credentials}` | `{error, documents, context}` | + +### Documents + +| Method | Path | Request body | Response | +|---|---|---|---| +| `POST` | `/api/get_document` | `GetDocumentPayload {uuid, credentials}` | `{error, document}` | +| `POST` | `/api/get_all_documents` | `SearchQueryPayload {query, labels, page, pageSize, credentials}` | `{documents, labels, error, totalDocuments}` | +| `POST` | `/api/delete_document` | `GetDocumentPayload {uuid, credentials}` | `{}` | +| `POST` | `/api/get_datacount` | `DatacountPayload {embedding_model, documentFilter, credentials}` | `{datacount}` | +| `POST` | `/api/get_labels` | `Credentials` | `{labels}` | +| `POST` | `/api/get_content` | `GetContentPayload {uuid, page, chunkScores, credentials}` | `{error, content, maxPage}` | +| `POST` | `/api/get_chunks` | `ChunksPayload {uuid, page, pageSize, credentials}` | `{error, chunks}` | +| `POST` | `/api/get_chunk` | `GetChunkPayload {uuid, embedder, credentials}` | `{error, chunk}` | +| `POST` | `/api/get_vectors` | `GetVectorPayload {uuid, showAll, credentials}` | `{error, vector_groups}` | + +### Suggestions + +| Method | Path | Request body | Response | +|---|---|---|---| +| `POST` | `/api/get_suggestions` | `GetSuggestionsPayload {query, limit, credentials}` | `{suggestions}` | +| `POST` | `/api/get_all_suggestions` | `GetAllSuggestionsPayload {page, pageSize, credentials}` | `{suggestions}` | +| `POST` | `/api/delete_suggestion` | `DeleteSuggestionPayload {uuid, credentials}` | `{}` | + +### Admin + +| Method | Path | Request body | Response | +|---|---|---|---| +| `POST` | `/api/reset` | `ResetPayload {resetMode, credentials}` | `{}` (resetMode: `"ALL"` \| `"DOCUMENTS"` \| `"CONFIG"` \| `"SUGGESTIONS"`) | +| `POST` | `/api/get_meta` | `Credentials` | `{error, node_payload, collection_payload}` | + +### WebSockets + +| Path | Direction | Description | +|---|---|---| +| `/ws/generate_stream` | Bidirectional, persistent | Send `GeneratePayload`, receive token stream | +| `/ws/import_files` | Bidirectional, persistent | Send `DataBatchPayload` chunks, receive `StatusReport` frames | + +### Static Assets + +| Method | Path | Description | +|---|---|---| +| `GET` | `/` | Serves `frontend/out/index.html` | +| `GET` | `/static/_next/*` | Next.js JS/CSS bundles | +| `GET` | `/static/*` | Other static files from the Next.js build | + +--- + +## Quick-Start Checklist for a New Generator + +``` +[ ] Create goldenverba/components/generation/MyGenerator.py +[ ] Class inherits from Generator (goldenverba/components/interfaces.py) +[ ] Set self.name to a unique string +[ ] Set self.description +[ ] Optionally set self.context_window +[ ] Add InputConfig entries to self.config (Model dropdown, API Key password if needed) +[ ] Implement async generate_stream(self, config, query, context, conversation) + - Must be an async generator (yield dicts) + - Yield {"message": str, "finish_reason": None | "stop"} + - Final yield must have finish_reason == "stop" +[ ] Implement prepare_messages() if needed +[ ] Use get_token() in __init__ and get_environment() in generate_stream +[ ] Register: import + add instance to generators list in generator_manager.py +[ ] If local-only, add self.name to the production exclusion set in generator_manager.py +[ ] Document env vars in PR description +[ ] Run: verba start and verify the generator appears in the UI dropdown +[ ] Run: pytest goldenverba/tests +[ ] Run: black goldenverba/ +``` diff --git a/goldenverba/components/chunk.py b/goldenverba/components/chunk.py index c876e0c750..9c7584c4dc 100644 --- a/goldenverba/components/chunk.py +++ b/goldenverba/components/chunk.py @@ -1,6 +1,3 @@ -from spacy.tokens import Doc, Span - - class Chunk: def __init__( self, diff --git a/goldenverba/components/chunking/MarkdownChunker.py b/goldenverba/components/chunking/MarkdownChunker.py index 058d6f521a..1101f36ae7 100644 --- a/goldenverba/components/chunking/MarkdownChunker.py +++ b/goldenverba/components/chunking/MarkdownChunker.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import contextlib with contextlib.suppress(Exception): diff --git a/goldenverba/components/chunking/chunker_manager.py b/goldenverba/components/chunking/chunker_manager.py new file mode 100644 index 0000000000..3fc0276354 --- /dev/null +++ b/goldenverba/components/chunking/chunker_manager.py @@ -0,0 +1,99 @@ +""" +chunker_manager.py +================== +Chunker component registry and ChunkerManager. + +To add a new Chunker: + 1. Implement it in this directory (goldenverba/components/chunking/) + 2. Import it below and add an instance to the `chunkers` list +""" + +import asyncio + +from goldenverba.components.document import Document +from goldenverba.components.interfaces import Chunker, Embedding +from goldenverba.server.helpers import LoggerManager +from goldenverba.server.types import FileConfig, FileStatus + +from goldenverba.components.chunking.TokenChunker import TokenChunker +from goldenverba.components.chunking.SentenceChunker import SentenceChunker +from goldenverba.components.chunking.RecursiveChunker import RecursiveChunker +from goldenverba.components.chunking.HTMLChunker import HTMLChunker +from goldenverba.components.chunking.MarkdownChunker import MarkdownChunker +from goldenverba.components.chunking.CodeChunker import CodeChunker +from goldenverba.components.chunking.JSONChunker import JSONChunker +from goldenverba.components.chunking.SemanticChunker import SemanticChunker + +# All available chunkers — add new instances here +chunkers = [ + TokenChunker(), + SentenceChunker(), + RecursiveChunker(), + SemanticChunker(), + HTMLChunker(), + MarkdownChunker(), + CodeChunker(), + JSONChunker(), +] + + +class ChunkerManager: + """Dispatches chunk() calls to the correct Chunker implementation.""" + + def __init__(self): + self.chunkers: dict[str, Chunker] = { + chunker.name: chunker for chunker in chunkers + } + + async def chunk( + self, + chunker: str, + fileConfig: FileConfig, + documents: list[Document], + embedder: Embedding, + logger: LoggerManager, + ) -> list[Document]: + try: + loop = asyncio.get_running_loop() + start_time = loop.time() + if chunker in self.chunkers: + config = fileConfig.rag_config["Chunker"].components[chunker].config + embedder_config = ( + fileConfig.rag_config["Embedder"].components[embedder.name].config + ) + chunked_documents = await self.chunkers[chunker].chunk( + config=config, + documents=documents, + embedder=embedder, + embedder_config=embedder_config, + ) + for chunked_document in chunked_documents: + chunked_document.meta["Chunker"] = ( + fileConfig.rag_config["Chunker"] + .components[chunker] + .model_dump() + ) + elapsed_time = round(loop.time() - start_time, 2) + if len(documents) == 1: + await logger.send_report( + fileConfig.fileID, + FileStatus.CHUNKING, + f"Split {fileConfig.filename} into {len(chunked_documents[0].chunks)} chunks", + took=elapsed_time, + ) + else: + await logger.send_report( + fileConfig.fileID, + FileStatus.CHUNKING, + f"Chunked all {len(chunked_documents)} documents with a total of {sum([len(document.chunks) for document in chunked_documents])} chunks", + took=elapsed_time, + ) + + await logger.send_report( + fileConfig.fileID, FileStatus.EMBEDDING, "", took=0 + ) + return chunked_documents + else: + raise Exception(f"{chunker} Chunker not found") + except Exception as e: + raise e diff --git a/goldenverba/components/client_manager.py b/goldenverba/components/client_manager.py new file mode 100644 index 0000000000..c6636ddbb1 --- /dev/null +++ b/goldenverba/components/client_manager.py @@ -0,0 +1,131 @@ +""" +client_manager.py +================= +Connection pool for Weaviate clients, isolated from VerbaManager so each module +stays focused on a single responsibility. + + ClientManager — maps hashed credentials → live WeaviateAsyncClient. + Concurrent requests with the same credentials share one connection. + A per-credential asyncio.Lock prevents duplicate connections + from opening in a race condition. + +Stale connections (idle > max_time minutes, or no longer responsive) are evicted +by clean_up(), which the /api/health endpoint triggers periodically. +""" + +import os +import asyncio +import hashlib +from datetime import datetime + +from wasabi import msg +from weaviate.client import WeaviateAsyncClient + +from goldenverba.server.types import Credentials +from goldenverba.components.verba_manager import VerbaManager + + +class ClientManager: + """ + Connection pool for Weaviate clients. + + Each unique set of credentials (deployment + URL + API key) maps to one + live WeaviateAsyncClient. The internal VerbaManager instance handles + the actual connection logic and all pipeline operations. + """ + + def __init__(self) -> None: + # {cred_hash: {"client": WeaviateAsyncClient, "timestamp": datetime}} + self.clients: dict[str, dict] = {} + self.manager: VerbaManager = VerbaManager() + self.max_time: int = 10 # minutes before an idle client is evicted + self.locks: dict[str, asyncio.Lock] = {} + + def hash_credentials(self, credentials: Credentials) -> str: + """Stable cache key derived from credentials; never stored or logged.""" + cred_string = f"{credentials.deployment}:{credentials.url}:{credentials.key}" + return hashlib.sha256(cred_string.encode()).hexdigest() + + def get_or_create_lock(self, cred_hash: str) -> asyncio.Lock: + """ + Return the per-credential lock, creating it if needed. + dict.setdefault() is atomic for dict operations, preventing a race where + two coroutines both see the key absent and create separate locks. + """ + self.locks.setdefault(cred_hash, asyncio.Lock()) + return self.locks[cred_hash] + + def heartbeat(self): + """Log the current connected-client count (debug aid).""" + msg.info(f"{len(self.clients)} clients connected") + for cred_hash, client in self.clients.items(): + msg.info(f"Client {cred_hash} connected at {client['timestamp']}") + + async def connect( + self, credentials: Credentials, port: str = "8080" + ) -> WeaviateAsyncClient: + """ + Return a live WeaviateAsyncClient for the given credentials. + + Returns an existing cached client if one exists; otherwise opens a new + connection under the per-credential lock to prevent duplicates. + Falls back to WEAVIATE_URL_VERBA / WEAVIATE_API_KEY_VERBA env vars when + the caller passes empty credentials (default deployment mode). + """ + # Work on a copy so we never mutate the caller's object. + _credentials = credentials.model_copy() + + if not _credentials.url and not _credentials.key: + _credentials.url = os.environ.get("WEAVIATE_URL_VERBA", "") + _credentials.key = os.environ.get("WEAVIATE_API_KEY_VERBA", "") + + cred_hash = self.hash_credentials(_credentials) + + lock = self.get_or_create_lock(cred_hash) + async with lock: + if cred_hash in self.clients: + msg.info("Found existing Client") + return self.clients[cred_hash]["client"] + else: + msg.warn("Connecting new Client") + client = await self.manager.connect(_credentials, port) + self.clients[cred_hash] = { + "client": client, + "timestamp": datetime.now(), + } + return client + + async def disconnect(self): + """Gracefully close all cached connections.""" + msg.warn("Disconnecting Clients!") + # Snapshot keys to avoid mutating the dict during iteration. + for cred_hash in list(self.clients.keys()): + await self.manager.disconnect(self.clients[cred_hash]["client"]) + + async def clean_up(self): + """ + Evict stale clients: those idle longer than max_time minutes, or whose + Weaviate connection is no longer healthy. Called by the /api/health endpoint. + """ + msg.info("Cleaning Clients Cache") + current_time = datetime.now() + clients_to_remove = [] + + # Snapshot to avoid RuntimeError if the dict changes concurrently. + for cred_hash, client_data in list(self.clients.items()): + time_difference = current_time - client_data["timestamp"] + if time_difference.total_seconds() / 60 > self.max_time: + clients_to_remove.append(cred_hash) + continue + client: WeaviateAsyncClient = client_data["client"] + if not await client.is_ready(): + clients_to_remove.append(cred_hash) + + for cred_hash in clients_to_remove: + if cred_hash in self.clients: + await self.manager.disconnect(self.clients[cred_hash]["client"]) + del self.clients[cred_hash] + msg.warn(f"Removed client: {cred_hash}") + + msg.info(f"Cleaned up {len(clients_to_remove)} clients") + self.heartbeat() diff --git a/goldenverba/components/document.py b/goldenverba/components/document.py index abb7fae886..692659b78b 100644 --- a/goldenverba/components/document.py +++ b/goldenverba/components/document.py @@ -1,10 +1,8 @@ from goldenverba.server.types import FileConfig from goldenverba.components.chunk import Chunk from spacy.tokens import Doc -from spacy.language import Language import spacy import json - from langdetect import detect @@ -39,7 +37,7 @@ def detect_language(text: str) -> str: elif detected_lang == "zh-tw" or detected_lang == "zh-hk": return "zh-hant" return detected_lang - except: + except Exception: return "unknown" diff --git a/goldenverba/components/embedding/LMStudioEmbedder.py b/goldenverba/components/embedding/LMStudioEmbedder.py new file mode 100644 index 0000000000..7705b3cd65 --- /dev/null +++ b/goldenverba/components/embedding/LMStudioEmbedder.py @@ -0,0 +1,128 @@ +import os +from typing import List + +import aiohttp +from wasabi import msg + +from goldenverba.components.interfaces import Embedding +from goldenverba.components.types import InputConfig +from goldenverba.components.util import get_environment, get_token + + +class LMStudioEmbedder(Embedding): + """LM Studio Embedder for Verba - Compatible with LM Studio's OpenAI-compatible API.""" + + def __init__(self): + super().__init__() + self.name = "LM Studio" + self.description = "Vectorizes documents and queries using LM Studio's locally hosted embedding models" + + # Default LM Studio configuration + api_key = get_token("LMSTUDIO_API_KEY") or "lm-studio" + base_url = os.getenv("LMSTUDIO_BASE_URL", "http://localhost:1234/v1") + + # Fetch available models + models = self.get_models(api_key, base_url) + default_model = os.getenv("LMSTUDIO_EMBED_MODEL", models[0] if models else "local-embedding") + + # Set up configuration + self.config = { + "Model": InputConfig( + type="dropdown", + value=default_model, + description="Select a LM Studio Embedding Model", + values=models, + ), + "URL": InputConfig( + type="text", + value=base_url, + description="LM Studio API Base URL (default: http://localhost:1234/v1)", + values=[], + ) + } + + # Add API Key config if not set in environment + if get_token("LMSTUDIO_API_KEY") is None: + self.config["API Key"] = InputConfig( + type="password", + value="lm-studio", + description="LM Studio API Key (often not required, default: 'lm-studio')", + values=[], + ) + + async def vectorize(self, config: dict, content: List[str]) -> List[List[float]]: + """Vectorize the input content using LM Studio's API.""" + model = config.get("Model", {"value": "local-embedding"}).value + api_key = get_environment( + config, "API Key", "LMSTUDIO_API_KEY", "lm-studio" + ) + base_url = get_environment( + config, "URL", "LMSTUDIO_BASE_URL", "http://localhost:1234/v1" + ) + + headers = { + "Authorization": f"Bearer {api_key}", + } + payload = {"input": content, "model": model} + + async with aiohttp.ClientSession() as session: + try: + async with session.post( + f"{base_url}/embeddings", + headers=headers, + json=payload, + timeout=aiohttp.ClientTimeout(total=60), + ) as response: + if response.status != 200: + error_text = await response.text() + raise Exception(f"LM Studio API error {response.status}: {error_text}") + + data = await response.json() + + if "data" not in data: + raise ValueError(f"Unexpected API response: {data}") + + embeddings = [item["embedding"] for item in data["data"]] + if len(embeddings) != len(content): + raise ValueError( + f"Mismatch in embedding count: got {len(embeddings)}, expected {len(content)}" + ) + + return embeddings + + except aiohttp.ClientError as e: + if isinstance(e, aiohttp.ClientResponseError) and e.status == 429: + raise Exception("Rate limit exceeded. Waiting before retrying...") + raise Exception(f"LM Studio API request failed: {str(e)}") + + except Exception as e: + msg.fail(f"LM Studio embedding error: {type(e).__name__} - {str(e)}") + raise + + @staticmethod + def get_models(token: str, url: str) -> List[str]: + """Fetch available embedding models from LM Studio API.""" + default_models = ["local-embedding"] + try: + import requests + + headers = {"Authorization": f"Bearer {token}"} + response = requests.get(f"{url}/models", headers=headers, timeout=10) + + if response.status_code == 200: + models_data = response.json() + if "data" in models_data: + # Get all models - LM Studio may not differentiate embedding vs chat models in the API + available_models = [model["id"] for model in models_data["data"]] + # Filter for embedding models if possible (some models have 'embed' in the name) + embedding_models = [m for m in available_models if 'embed' in m.lower()] + + # Return embedding models if found, otherwise return all models + return embedding_models if embedding_models else available_models + + msg.info("Could not fetch embedding models from LM Studio, using default") + return default_models + + except Exception as e: + msg.info(f"Failed to fetch LM Studio embedding models: {str(e)}") + return default_models diff --git a/goldenverba/components/embedding/SentenceTransformersEmbedder.py b/goldenverba/components/embedding/SentenceTransformersEmbedder.py index 406f99e341..aa4591522d 100644 --- a/goldenverba/components/embedding/SentenceTransformersEmbedder.py +++ b/goldenverba/components/embedding/SentenceTransformersEmbedder.py @@ -1,10 +1,12 @@ +import asyncio + from goldenverba.components.interfaces import Embedding from goldenverba.components.types import InputConfig try: from sentence_transformers import SentenceTransformer -except Exception as e: - pass +except ImportError: + SentenceTransformer = None class SentenceTransformersEmbedder(Embedding): @@ -32,12 +34,26 @@ def __init__(self): ], ), } + # Cache loaded models by name to avoid reloading from disk on every call + self._model_cache: dict = {} + + def _get_model(self, model_name: str): + if SentenceTransformer is None: + raise ImportError( + "sentence_transformers is not installed. " + "Install it with: pip install goldenverba[huggingface]" + ) + if model_name not in self._model_cache: + self._model_cache[model_name] = SentenceTransformer(model_name) + return self._model_cache[model_name] - async def vectorize(self, config: dict, content: list[str]) -> list[float]: + async def vectorize(self, config: dict, content: list[str]) -> list[list[float]]: try: model_name = config.get("Model").value - model = SentenceTransformer(model_name) - embeddings = model.encode(content).tolist() - return embeddings + model = self._get_model(model_name) + # model.encode() is synchronous and CPU-bound — run in thread pool + # to avoid blocking the async event loop + embeddings = await asyncio.to_thread(model.encode, content) + return embeddings.tolist() except Exception as e: raise Exception(f"Failed to vectorize chunks: {str(e)}") diff --git a/goldenverba/components/embedding/embedding_manager.py b/goldenverba/components/embedding/embedding_manager.py new file mode 100644 index 0000000000..595e8eea9c --- /dev/null +++ b/goldenverba/components/embedding/embedding_manager.py @@ -0,0 +1,173 @@ +""" +embedding_manager.py +==================== +Embedder component registry and EmbeddingManager. + +To add a new Embedder: + 1. Implement it in this directory (goldenverba/components/embedding/) + 2. Import it below and add an instance to the embedders list + +VERBA_PRODUCTION env var +------------------------ +When set to "Production", local-only embedders (Ollama, SentenceTransformers) are +excluded. This is used for hosted deployments where those services aren't available. +""" + +import os +import asyncio + +from wasabi import msg +from sklearn.decomposition import PCA + +from goldenverba.components.document import Document +from goldenverba.components.interfaces import Embedding +from goldenverba.server.helpers import LoggerManager +from goldenverba.server.types import FileConfig, FileStatus + +from goldenverba.components.embedding.OpenAIEmbedder import OpenAIEmbedder +from goldenverba.components.embedding.CohereEmbedder import CohereEmbedder +from goldenverba.components.embedding.OllamaEmbedder import OllamaEmbedder +from goldenverba.components.embedding.UpstageEmbedder import UpstageEmbedder +from goldenverba.components.embedding.WeaviateEmbedder import WeaviateEmbedder +from goldenverba.components.embedding.VoyageAIEmbedder import VoyageAIEmbedder +from goldenverba.components.embedding.SentenceTransformersEmbedder import ( + SentenceTransformersEmbedder, +) +from goldenverba.components.embedding.LMStudioEmbedder import LMStudioEmbedder + +# Local-only embedders are excluded in Production mode (hosted deployments) +_is_production = os.getenv("VERBA_PRODUCTION") == "Production" + +embedders = [ + e for e in [ + OllamaEmbedder(), + SentenceTransformersEmbedder(), + WeaviateEmbedder(), + UpstageEmbedder(), + VoyageAIEmbedder(), + CohereEmbedder(), + OpenAIEmbedder(), + LMStudioEmbedder(), + ] + # OllamaEmbedder and SentenceTransformersEmbedder require local services + if not _is_production or e.name not in {"Ollama", "SentenceTransformers"} +] + + +class EmbeddingManager: + """Dispatches vectorize() calls to the correct Embedder implementation.""" + + def __init__(self): + self.embedders: dict[str, Embedding] = { + embedder.name: embedder for embedder in embedders + } + + async def vectorize( + self, + embedder: str, + fileConfig: FileConfig, + documents: list[Document], + logger: LoggerManager, + ) -> list[Document]: + """Vectorizes chunks in batches + @parameter: documents : Document - Verba document + @returns Document - Document with vectorized chunks + """ + try: + loop = asyncio.get_running_loop() + start_time = loop.time() + if embedder in self.embedders: + config = fileConfig.rag_config["Embedder"].components[embedder].config + + for document in documents: + content = [ + document.metadata + "\n" + chunk.content + for chunk in document.chunks + ] + embeddings = await self.batch_vectorize(embedder, config, content) + + if len(embeddings) >= 3: + pca = PCA(n_components=3) + generated_pca_embeddings = pca.fit_transform(embeddings) + pca_embeddings = [ + pca_.tolist() for pca_ in generated_pca_embeddings + ] + else: + pca_embeddings = [embedding[0:3] for embedding in embeddings] + + for vector, chunk, pca_ in zip( + embeddings, document.chunks, pca_embeddings + ): + chunk.vector = vector + chunk.pca = pca_ + + document.meta["Embedder"] = ( + fileConfig.rag_config["Embedder"] + .components[embedder] + .model_dump() + ) + + elapsed_time = round(loop.time() - start_time, 2) + await logger.send_report( + fileConfig.fileID, + FileStatus.EMBEDDING, + f"Vectorized all chunks", + took=elapsed_time, + ) + await logger.send_report( + fileConfig.fileID, FileStatus.INGESTING, "", took=0 + ) + return documents + else: + raise Exception(f"{embedder} Embedder not found") + except Exception as e: + raise e + + async def batch_vectorize( + self, embedder: str, config: dict, content: list[str] + ) -> list[list[float]]: + """Vectorize content in batches""" + try: + batches = [ + content[i : i + self.embedders[embedder].max_batch_size] + for i in range(0, len(content), self.embedders[embedder].max_batch_size) + ] + msg.info(f"Vectorizing {len(content)} chunks in {len(batches)} batches") + tasks = [ + self.embedders[embedder].vectorize(config, batch) for batch in batches + ] + results = await asyncio.gather(*tasks, return_exceptions=True) + + # Check if all tasks were successful + errors = [r for r in results if isinstance(r, Exception)] + if errors: + error_messages = [str(e) for e in errors] + raise Exception( + f"Vectorization failed for some batches: {', '.join(error_messages)}" + ) + + # Flatten the results + flattened_results = [item for sublist in results for item in sublist] + + # Verify the number of vectors matches the input content + if len(flattened_results) != len(content): + raise Exception( + f"Mismatch in vectorization results: expected {len(content)} vectors, got {len(flattened_results)}" + ) + + return flattened_results + except Exception as e: + raise Exception(f"Batch vectorization failed: {str(e)}") + + async def vectorize_query( + self, embedder: str, content: str, rag_config: dict + ) -> list[float]: + try: + if embedder in self.embedders: + config = rag_config["Embedder"].components[embedder].config + embeddings = await self.embedders[embedder].vectorize(config, [content]) + return embeddings[0] + else: + raise Exception(f"{embedder} Embedder not found") + except Exception as e: + raise e diff --git a/goldenverba/components/generation/AnthrophicGenerator.py b/goldenverba/components/generation/AnthrophicGenerator.py index f49b8eab89..5f7348ea58 100644 --- a/goldenverba/components/generation/AnthrophicGenerator.py +++ b/goldenverba/components/generation/AnthrophicGenerator.py @@ -72,6 +72,7 @@ async def generate_stream( self.url, json=data, headers=headers, + timeout=aiohttp.ClientTimeout(connect=10, total=300), ) as response: if response.status != 200: error_json = await response.json() diff --git a/goldenverba/components/generation/CohereGenerator.py b/goldenverba/components/generation/CohereGenerator.py index 77345de956..f9518254bc 100644 --- a/goldenverba/components/generation/CohereGenerator.py +++ b/goldenverba/components/generation/CohereGenerator.py @@ -76,7 +76,10 @@ async def generate_stream( try: async with aiohttp.ClientSession() as session: async with session.post( - self.url + "/chat", json=data, headers=headers + self.url + "/chat", + json=data, + headers=headers, + timeout=aiohttp.ClientTimeout(connect=10, total=300), ) as response: if response.status == 200: async for line in response.content: diff --git a/goldenverba/components/generation/DeepSeekGenerator.py b/goldenverba/components/generation/DeepSeekGenerator.py new file mode 100644 index 0000000000..b18ee3bba9 --- /dev/null +++ b/goldenverba/components/generation/DeepSeekGenerator.py @@ -0,0 +1,190 @@ +import os +from dotenv import load_dotenv +from goldenverba.components.interfaces import Generator +from goldenverba.components.types import InputConfig +from goldenverba.components.util import get_environment, get_token +from typing import List +import httpx +import json +from wasabi import msg + +load_dotenv() + +DEEPSEEK_BASE_URL = "https://api.deepseek.com/v1" +DEFAULT_MODEL_LIST = ["deepseek-chat", "deepseek-reasoner"] + + +class DeepSeekGenerator(Generator): + """ + DeepSeek Generator with reasoning model support. + Supports both deepseek-chat and deepseek-reasoner (R1), + including optional display of the model's thinking process. + """ + + def __init__(self): + super().__init__() + self.name = "DeepSeek" + self.description = "Using DeepSeek models to generate answers, with support for reasoning models (R1)" + self.context_window = 10000 + + api_key = get_token("DEEPSEEK_API_KEY") + base_url = os.getenv("DEEPSEEK_BASE_URL", DEEPSEEK_BASE_URL) + models = self.get_models(api_key, base_url) + default_model = os.getenv("DEEPSEEK_MODEL", models[0]) + + self.config["Model"] = InputConfig( + type="dropdown", + value=default_model, + description="Select a DeepSeek Model", + values=models, + ) + + self.config["Show Reasoning"] = InputConfig( + type="bool", + value=False, + description="Show the model's thinking process (for reasoning models like deepseek-reasoner)", + values=[], + ) + + if api_key is None: + self.config["API Key"] = InputConfig( + type="password", + value="", + description="You can set your DeepSeek API Key here or set it as environment variable `DEEPSEEK_API_KEY`", + values=[], + ) + if os.getenv("DEEPSEEK_BASE_URL") is None: + self.config["URL"] = InputConfig( + type="text", + value=DEEPSEEK_BASE_URL, + description="You can change the Base URL here if needed", + values=[], + ) + + async def generate_stream( + self, + config: dict, + query: str, + context: str, + conversation: list[dict] = [], + ): + system_message = config.get("System Message").value + model = config.get("Model", {"value": "deepseek-chat"}).value + show_reasoning = config.get("Show Reasoning", {"value": False}).value + api_key = get_environment( + config, "API Key", "DEEPSEEK_API_KEY", "No DeepSeek API Key found" + ) + api_url = get_environment( + config, "URL", "DEEPSEEK_BASE_URL", DEEPSEEK_BASE_URL + ) + + messages = self.prepare_messages(query, context, conversation, system_message) + + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + } + data = { + "messages": messages, + "model": model, + "stream": True, + } + + in_thinking = False + + async with httpx.AsyncClient() as client: + async with client.stream( + "POST", + f"{api_url}/chat/completions", + json=data, + headers=headers, + timeout=httpx.Timeout(connect=10, read=300), + ) as response: + async for line in response.aiter_lines(): + if line.startswith("data: "): + if line.strip() == "data: [DONE]": + break + try: + json_line = json.loads(line[6:]) + except json.JSONDecodeError: + continue + choice = json_line["choices"][0] + + # Handle reasoning_content from deepseek-reasoner + if "delta" in choice: + delta = choice["delta"] + reasoning = delta.get("reasoning_content") + content = delta.get("content") + + if reasoning and show_reasoning: + if not in_thinking: + in_thinking = True + yield { + "message": "\n<details><summary>💭 Reasoning</summary>\n\n", + "finish_reason": None, + } + yield { + "message": reasoning, + "finish_reason": None, + } + + if content: + if in_thinking: + in_thinking = False + yield { + "message": "\n</details>\n\n", + "finish_reason": None, + } + yield { + "message": content, + "finish_reason": choice.get("finish_reason"), + } + + if choice.get("finish_reason") == "stop": + if in_thinking: + yield { + "message": "\n</details>\n\n", + "finish_reason": None, + } + yield { + "message": "", + "finish_reason": "stop", + } + + def prepare_messages( + self, query: str, context: str, conversation: list[dict], system_message: str + ) -> list[dict]: + messages = [ + { + "role": "system", + "content": system_message, + } + ] + + for message in conversation: + messages.append({"role": message.type, "content": message.content}) + + messages.append( + { + "role": "user", + "content": f"Answer this query: '{query}' with this provided context: {context}", + } + ) + + return messages + + def get_models(self, token: str, url: str) -> List[str]: + """Fetch available models from DeepSeek API.""" + try: + if token is None: + return DEFAULT_MODEL_LIST + + import requests + + headers = {"Authorization": f"Bearer {token}"} + response = requests.get(f"{url}/models", headers=headers, timeout=10) + response.raise_for_status() + return [model["id"] for model in response.json()["data"]] + except Exception as e: + msg.info(f"Failed to fetch DeepSeek models: {str(e)}") + return DEFAULT_MODEL_LIST diff --git a/goldenverba/components/generation/GeminiGenerator.py b/goldenverba/components/generation/GeminiGenerator.py index f3c8847cb4..0f4c4ef44e 100644 --- a/goldenverba/components/generation/GeminiGenerator.py +++ b/goldenverba/components/generation/GeminiGenerator.py @@ -3,7 +3,7 @@ try: import vertexai.preview from vertexai.preview.generative_models import GenerativeModel, Content, Part -except: +except ImportError: pass from wasabi import msg diff --git a/goldenverba/components/generation/GroqGenerator.py b/goldenverba/components/generation/GroqGenerator.py index b1c8deef0c..983972b8c8 100644 --- a/goldenverba/components/generation/GroqGenerator.py +++ b/goldenverba/components/generation/GroqGenerator.py @@ -90,7 +90,10 @@ async def generate_stream( try: async with aiohttp.ClientSession() as session: async with session.post( - self.url + "/chat/completions", json=data, headers=headers + self.url + "/chat/completions", + json=data, + headers=headers, + timeout=aiohttp.ClientTimeout(connect=10, total=300), ) as response: if response.status == 200: async for line in response.content: diff --git a/goldenverba/components/generation/LMStudioGenerator.py b/goldenverba/components/generation/LMStudioGenerator.py new file mode 100644 index 0000000000..ee758f04b5 --- /dev/null +++ b/goldenverba/components/generation/LMStudioGenerator.py @@ -0,0 +1,170 @@ +import os +from dotenv import load_dotenv +from goldenverba.components.interfaces import Generator +from goldenverba.components.types import InputConfig +from goldenverba.components.util import get_environment, get_token +from typing import List +import httpx +import json +from wasabi import msg + +load_dotenv() + + +class LMStudioGenerator(Generator): + """ + LM Studio Generator - Compatible with LM Studio's OpenAI-compatible API. + """ + + def __init__(self): + super().__init__() + self.name = "LM Studio" + self.description = "Using LM Studio's locally hosted models via OpenAI-compatible API" + self.context_window = 10000 + + # Default LM Studio URL + base_url = os.getenv("LMSTUDIO_BASE_URL", "http://localhost:1234/v1") + api_key = get_token("LMSTUDIO_API_KEY") or "lm-studio" # LM Studio often doesn't need real API key + + models = self.get_models(api_key, base_url) + default_model = os.getenv("LMSTUDIO_MODEL", models[0] if models else "local-model") + + self.config["Model"] = InputConfig( + type="dropdown", + value=default_model, + description="Select a LM Studio Model", + values=models, + ) + + # Always show URL config for LM Studio since it's local + self.config["URL"] = InputConfig( + type="text", + value=base_url, + description="LM Studio API Base URL (default: http://localhost:1234/v1)", + values=[], + ) + + # API Key is optional for LM Studio + if get_token("LMSTUDIO_API_KEY") is None: + self.config["API Key"] = InputConfig( + type="password", + value="lm-studio", + description="LM Studio API Key (often not required, default: 'lm-studio')", + values=[], + ) + + async def generate_stream( + self, + config: dict, + query: str, + context: str, + conversation: list[dict] = [], + ): + system_message = config.get("System Message").value + model = config.get("Model", {"value": "local-model"}).value + lmstudio_key = get_environment( + config, "API Key", "LMSTUDIO_API_KEY", "lm-studio" + ) + lmstudio_url = get_environment( + config, "URL", "LMSTUDIO_BASE_URL", "http://localhost:1234/v1" + ) + + messages = self.prepare_messages(query, context, conversation, system_message) + + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {lmstudio_key}", + } + data = { + "messages": messages, + "model": model, + "stream": True, + "temperature": 0.7, + "max_tokens": 2048, + } + + try: + async with httpx.AsyncClient(timeout=httpx.Timeout(connect=10, read=300)) as client: + async with client.stream( + "POST", + f"{lmstudio_url}/chat/completions", + json=data, + headers=headers, + ) as response: + if response.status_code != 200: + error_text = await response.aread() + raise Exception(f"LM Studio API error {response.status_code}: {error_text}") + + async for line in response.aiter_lines(): + if line.startswith("data: "): + if line.strip() == "data: [DONE]": + yield { + "message": "", + "finish_reason": "stop", + } + break + try: + json_line = json.loads(line[6:]) + choice = json_line.get("choices", [{}])[0] + if "delta" in choice and "content" in choice["delta"]: + yield { + "message": choice["delta"]["content"], + "finish_reason": choice.get("finish_reason"), + } + elif "finish_reason" in choice: + yield { + "message": "", + "finish_reason": choice["finish_reason"], + } + except json.JSONDecodeError: + continue # Skip malformed lines + except Exception as e: + msg.fail(f"LM Studio generation failed: {str(e)}") + yield { + "message": f"Error: {str(e)}", + "finish_reason": "stop", + } + + def prepare_messages( + self, query: str, context: str, conversation: list[dict], system_message: str + ) -> list[dict]: + messages = [ + { + "role": "system", + "content": system_message, + } + ] + + for message in conversation: + messages.append({"role": message.type, "content": message.content}) + + messages.append( + { + "role": "user", + "content": f"Answer this query: '{query}' with this provided context: {context}", + } + ) + + return messages + + def get_models(self, token: str, url: str) -> List[str]: + """Fetch available models from LM Studio API.""" + default_models = ["local-model"] + try: + import requests + + headers = {"Authorization": f"Bearer {token}"} + response = requests.get(f"{url}/models", headers=headers, timeout=10) + + if response.status_code == 200: + models_data = response.json() + if "data" in models_data: + available_models = [model["id"] for model in models_data["data"]] + return available_models if available_models else default_models + + msg.info("Could not fetch models from LM Studio, using default") + return default_models + + except Exception as e: + msg.info(f"Failed to fetch LM Studio models: {str(e)}") + return default_models diff --git a/goldenverba/components/generation/NovitaGenerator.py b/goldenverba/components/generation/NovitaGenerator.py index cc12cde554..d0dff871a5 100644 --- a/goldenverba/components/generation/NovitaGenerator.py +++ b/goldenverba/components/generation/NovitaGenerator.py @@ -72,7 +72,7 @@ async def generate_stream( url=f"{novita_url}/chat/completions", json=data, headers=headers, - timeout=None, + timeout=aiohttp.ClientTimeout(connect=10, total=300), ) as response: if response.status == 200: async for line in response.content: diff --git a/goldenverba/components/generation/OllamaGenerator.py b/goldenverba/components/generation/OllamaGenerator.py index 800565c49b..6866ae7f59 100644 --- a/goldenverba/components/generation/OllamaGenerator.py +++ b/goldenverba/components/generation/OllamaGenerator.py @@ -47,7 +47,11 @@ async def generate_stream( try: async with aiohttp.ClientSession() as session: - async with session.post(urljoin(self.url, "/api/chat"), json=data) as response: + async with session.post( + urljoin(self.url, "/api/chat"), + json=data, + timeout=aiohttp.ClientTimeout(connect=10, total=300), + ) as response: async for line in response.content: if line.strip(): yield self._process_response(line) diff --git a/goldenverba/components/generation/OpenAIGenerator.py b/goldenverba/components/generation/OpenAIGenerator.py index 5064b24d1f..84c1185a99 100644 --- a/goldenverba/components/generation/OpenAIGenerator.py +++ b/goldenverba/components/generation/OpenAIGenerator.py @@ -83,7 +83,7 @@ async def generate_stream( f"{openai_url}/chat/completions", json=data, headers=headers, - timeout=None, + timeout=httpx.Timeout(connect=10, read=300), ) as response: async for line in response.aiter_lines(): if line.startswith("data: "): diff --git a/goldenverba/components/generation/UpstageGenerator.py b/goldenverba/components/generation/UpstageGenerator.py index ed9a367675..b261d3928e 100644 --- a/goldenverba/components/generation/UpstageGenerator.py +++ b/goldenverba/components/generation/UpstageGenerator.py @@ -87,7 +87,7 @@ async def generate_stream( f"{base_url}/chat/completions", json=data, headers=headers, - timeout=None, + timeout=httpx.Timeout(connect=10, read=300), ) as response: response.raise_for_status() async for line in response.aiter_lines(): diff --git a/goldenverba/components/generation/generator_manager.py b/goldenverba/components/generation/generator_manager.py new file mode 100644 index 0000000000..c38bd98093 --- /dev/null +++ b/goldenverba/components/generation/generator_manager.py @@ -0,0 +1,126 @@ +""" +generator_manager.py +==================== +Generator component registry and GeneratorManager. + +To add a new Generator: + 1. Implement it in this directory (goldenverba/components/generation/) + 2. Import it below and add an instance to the generators list + +VERBA_PRODUCTION env var +------------------------ +When set to "Production", local-only generators (Ollama, Groq, Novita AI) are +excluded. This is used for hosted deployments where those services aren't available. +""" + +import os + +from wasabi import msg + +from goldenverba.components.interfaces import Generator + +from goldenverba.components.generation.CohereGenerator import CohereGenerator +from goldenverba.components.generation.AnthrophicGenerator import AnthropicGenerator +from goldenverba.components.generation.OllamaGenerator import OllamaGenerator +from goldenverba.components.generation.OpenAIGenerator import OpenAIGenerator +from goldenverba.components.generation.GroqGenerator import GroqGenerator +from goldenverba.components.generation.NovitaGenerator import NovitaGenerator +from goldenverba.components.generation.UpstageGenerator import UpstageGenerator +from goldenverba.components.generation.DeepSeekGenerator import DeepSeekGenerator +from goldenverba.components.generation.LMStudioGenerator import LMStudioGenerator + +try: + import tiktoken +except Exception: + msg.warn("tiktoken not installed, your base installation might be corrupted.") + +# Local-only generators are excluded in Production mode (hosted deployments) +_is_production = os.getenv("VERBA_PRODUCTION") == "Production" + +generators = [ + g for g in [ + OllamaGenerator(), + OpenAIGenerator(), + AnthropicGenerator(), + CohereGenerator(), + GroqGenerator(), + NovitaGenerator(), + UpstageGenerator(), + DeepSeekGenerator(), + LMStudioGenerator(), + ] + # OllamaGenerator, GroqGenerator, NovitaGenerator excluded in Production + if not _is_production or g.name not in {"Ollama", "Groq", "Novita AI"} +] + + +class GeneratorManager: + """Dispatches generate_stream() calls to the correct Generator implementation.""" + + def __init__(self): + self.generators: dict[str, Generator] = { + generator.name: generator for generator in generators + } + + async def generate_stream(self, rag_config, query, context, conversation): + """Generate a stream of response dicts based on a list of queries and list of contexts, and includes conversational context + @parameter: queries : list[str] - List of queries + @parameter: context : list[str] - List of contexts + @parameter: conversation : dict - Conversational context + @returns Iterator[dict] - Token response generated by the Generator in this format {system:TOKEN, finish_reason:stop or empty}. + """ + + generator = rag_config["Generator"].selected + generator_config = ( + rag_config["Generator"].components[rag_config["Generator"].selected].config + ) + + if generator not in self.generators: + raise Exception(f"Generator {generator} not found") + + async for result in self.generators[generator].generate_stream( + generator_config, query, context, conversation + ): + yield result + + def truncate_conversation_dicts( + self, conversation_dicts: list[dict[str, any]], max_tokens: int + ) -> list[dict[str, any]]: + """ + Truncate a list of conversation dictionaries to fit within a specified maximum token limit. + + @parameter conversation_dicts: List[Dict[str, any]] - A list of conversation dictionaries that may contain various keys, where 'content' key is present and contains text data. + @parameter max_tokens: int - The maximum number of tokens that the combined content of the truncated conversation dictionaries should not exceed. + + @returns List[Dict[str, any]]: A list of conversation dictionaries that have been truncated so that their combined content respects the max_tokens limit. The list is returned in the original order of conversation with the most recent conversation being truncated last if necessary. + + """ + encoding = tiktoken.encoding_for_model("gpt-3.5-turbo") + accumulated_tokens = 0 + truncated_conversation_dicts = [] + + # Start with the newest conversations + for item_dict in reversed(conversation_dicts): + item_tokens = encoding.encode(item_dict["content"], disallowed_special=()) + + # If adding the entire new item exceeds the max tokens + if accumulated_tokens + len(item_tokens) > max_tokens: + # Calculate how many tokens we can add from this item + remaining_space = max_tokens - accumulated_tokens + truncated_content = encoding.decode(item_tokens[:remaining_space]) + + # Create a new truncated item dictionary + truncated_item_dict = { + "type": item_dict["type"], + "content": truncated_content, + "typewriter": item_dict["typewriter"], + } + + truncated_conversation_dicts.append(truncated_item_dict) + break + + truncated_conversation_dicts.append(item_dict) + accumulated_tokens += len(item_tokens) + + # The list has been built in reverse order so we reverse it again + return list(reversed(truncated_conversation_dicts)) diff --git a/goldenverba/components/reader/AssemblyAIAPI.py b/goldenverba/components/reader/AssemblyAIAPI.py deleted file mode 100644 index d34af1f8cb..0000000000 --- a/goldenverba/components/reader/AssemblyAIAPI.py +++ /dev/null @@ -1,137 +0,0 @@ -import base64 -import io -import os - -import requests -from wasabi import msg -import aiohttp -import assemblyai as aai - -from goldenverba.components.document import Document, create_document -from goldenverba.components.interfaces import Reader -from goldenverba.server.types import FileConfig -from goldenverba.components.util import get_environment -from goldenverba.components.types import InputConfig - - -class AssemblyAIReader(Reader): - """ - AssemblyAI API Reader for importing multiple file types using the AssemblyAI.com API. - """ - - def __init__(self): - super().__init__() - self.extension = [ - ".3ga", - ".webm", - ".8svx", - ".mts", - ".m2ts", - ".ts", - ".aac", - ".mov", - ".ac3", - ".mp2", - ".aif", - ".mp4", - ".m4p", - ".m4v", - ".aiff", - ".mxf", - ".alac", - ".amr", - ".ape", - ".au", - ".dss", - ".flac", - ".flv", - ".m4a", - ".m4b", - ".m4p", - ".m4r", - ".mp3", - ".mpga", - ".ogg", - ".oga", - ".mogg", - ".opus", - ".qcp", - ".tta", - ".voc", - ".wav", - ".wma", - ".wv", - ] - self.requires_env = ["ASSEMBLYAI_API_KEY"] - self.name = "AssemblyAI" - self.description = "Uses the AssemblyAI API to import multiple file types such as plain text and documents" - self.config = { - "Quality": InputConfig( - type="dropdown", - value="best", - description="Set the transcription quality", - values=["nano", "best"], - ) - } - - if os.getenv("ASSEMBLYAI_API_KEY") is None: - self.config["API Key"] = InputConfig( - type="password", - value="", - description="Set your AssemblyAI API Key here or set it as an environment variable `ASSEMBLYAI_API_KEY`", - values=[], - ) - - async def load( - self, config: dict[str, InputConfig], fileConfig: FileConfig - ) -> list[Document]: - """ - Load and process a file using the AssemblyAI API. - """ - # Validate and get API credentials - token = get_environment( - config, - "API Key", - "ASSEMBLYAI_API_KEY", - "No AssemblyAI API Key detected", - ) - aai.settings.api_key = token - - # Validate quality - quality = config["Quality"].value - if quality not in ["nano", "best"]: - raise ValueError(f"Invalid quality: {quality}") - - aaiConfig = aai.TranscriptionConfig(speech_model=aai.SpeechModel.nano) - if quality == "best": - aaiConfig = aai.TranscriptionConfig(speech_model=aai.SpeechModel.best) - - msg.info(f"Loading {fileConfig.filename}") - - file_data = aiohttp.FormData() - file_bytes = io.BytesIO(base64.b64decode(fileConfig.content)) - file_data.add_field( - "files", - file_bytes, - filename=f"{fileConfig.filename}.{fileConfig.extension}", - ) - - try: - transcriber = aai.Transcriber(config=aaiConfig) - transcript = transcriber.transcribe(file_bytes) - if transcript.error: - raise Exception( - f"AssemblyAI API failed to transcribe {fileConfig.filename}: {transcript.error}" - ) - if transcript.text is None: - raise Exception( - f"AssemblyAI API failed to transcribe {fileConfig.filename}, no text returned" - ) - return [create_document(transcript.text, fileConfig)] - - except requests.RequestException as e: - raise Exception( - f"AssemblyAI API request failed for {fileConfig.filename}: {str(e)}" - ) - except Exception as e: - raise Exception(f"Failed to process {fileConfig.filename}: {str(e)}") diff --git a/goldenverba/components/reader/BasicReader.py b/goldenverba/components/reader/BasicReader.py index fc5a2ddb5c..e933c67b26 100644 --- a/goldenverba/components/reader/BasicReader.py +++ b/goldenverba/components/reader/BasicReader.py @@ -165,7 +165,10 @@ async def load_pdf_file(self, decoded_bytes: bytes) -> str: raise ImportError("pypdf is not installed. Cannot process PDF files.") pdf_bytes = io.BytesIO(decoded_bytes) reader = PdfReader(pdf_bytes) - return "\n\n".join(page.extract_text() for page in reader.pages) + # extract_text() returns None on image-only or corrupt pages; filter those out + return "\n\n".join( + text for page in reader.pages if (text := page.extract_text()) + ) async def load_docx_file(self, decoded_bytes: bytes) -> str: """Load and extract text from a DOCX file.""" diff --git a/goldenverba/components/reader/FirecrawlReader.py b/goldenverba/components/reader/FirecrawlReader.py deleted file mode 100644 index 80e82cb75a..0000000000 --- a/goldenverba/components/reader/FirecrawlReader.py +++ /dev/null @@ -1,237 +0,0 @@ -import base64 -import aiohttp -import asyncio -import os -from typing import List, Tuple - -from wasabi import msg - -from goldenverba.components.document import Document -from goldenverba.components.interfaces import Reader -from goldenverba.server.types import FileConfig -from goldenverba.components.reader.BasicReader import BasicReader -from goldenverba.components.util import get_environment -from goldenverba.components.types import InputConfig - - -class FirecrawlReader(Reader): - """ - FirecrawlReader uses the Firecrawl API to scrape or crawl websites and ingest them into Verba. - """ - - def __init__(self): - super().__init__() - self.name = "Firecrawl" - self.type = "URL" - self.description = "Use Firecrawl to scrape websites and ingest them into Verba" - self.config = { - "Mode": InputConfig( - type="dropdown", - value="Scrape", - description="Switch between scraping and crawling. Note that crawling can take some time.", - values=["Crawl", "Scrape"], - ), - "URLs": InputConfig( - type="multi", - value="", - description="Add URLs to retrieve data from", - values=[], - ), - } - - if os.getenv("FIRECRAWL_API_KEY") is None: - self.config["Firecrawl API Key"] = InputConfig( - type="password", - value="", - description="You can set your Firecrawl API Key or set it as environment variable `FIRECRAWL_API_KEY`", - values=[], - ) - - async def load(self, config: dict, fileConfig: FileConfig) -> List[Document]: - """ - Load documents from URLs using Firecrawl API. - """ - reader = BasicReader() - urls = config["URLs"].values - mode = config["Mode"].value - token = get_environment( - config, - "Firecrawl API Key", - "FIRECRAWL_API_KEY", - "No Firecrawl API Key detected", - ) - - raw_documents = await self.firecrawl(mode, urls, token) - documents = [] - - for title, content, source_url in raw_documents: - content_bytes = content.encode("utf-8") - base64_content = base64.b64encode(content_bytes).decode("utf-8") - - new_file_config = FileConfig( - fileID=fileConfig.fileID, - filename=title, - isURL=False, - overwrite=fileConfig.overwrite, - extension="md", - source=source_url, - content=base64_content, - labels=fileConfig.labels, - rag_config=fileConfig.rag_config, - file_size=len(content_bytes), - status=fileConfig.status, - status_report=fileConfig.status_report, - metadata=fileConfig.metadata, - ) - document = await reader.load(config, new_file_config) - documents.append(document[0]) - - return documents - - async def handle_response(self, response: aiohttp.ClientResponse) -> dict: - """ - Handle the API response and raise an exception if the status is not 200. - """ - if response.status != 200: - text = await response.text() - raise Exception(f"Firecrawl Error: {response.status}, {text}") - return await response.json() - - async def firecrawl( - self, mode: str, urls: List[str], token: str - ) -> List[Tuple[str, str, str]]: - """ - Perform scraping or crawling using Firecrawl API. - """ - crawl_url = "https://api.firecrawl.dev/v0/crawl" - scrape_url = "https://api.firecrawl.dev/v0/scrape" - documents = [] - - headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {token}", - } - - async with aiohttp.ClientSession() as session: - tasks = [] - for url in urls: - request_data = {"url": url} - if mode == "Scrape": - task = self.scrape_url(session, scrape_url, headers, request_data) - else: - task = self.handle_crawl(session, crawl_url, headers, request_data) - tasks.append(task) - - results = await asyncio.gather(*tasks, return_exceptions=True) - for result in results: - if isinstance(result, Exception): - msg.warn(f"Failed to process URL: {str(result)}") - else: - documents.extend(result) - - if not documents: - raise Exception( - "Firecrawl was not able to load any documents, please check your API Key and settings" - ) - - return documents - - async def scrape_url( - self, - session: aiohttp.ClientSession, - scrape_url: str, - headers: dict, - request_data: dict, - ) -> List[Tuple[str, str, str]]: - """ - Scrape a single URL using Firecrawl API. - """ - async with session.post( - scrape_url, headers=headers, json=request_data - ) as response: - response_data = await self.handle_response(response) - if "data" in response_data and response_data.get("success", False): - return [ - ( - response_data["data"]["metadata"]["title"], - response_data["data"]["markdown"], - request_data["url"], - ) - ] - return [] - - async def handle_crawl( - self, - session: aiohttp.ClientSession, - crawl_url: str, - headers: dict, - request_data: dict, - ) -> List[Tuple[str, str, str]]: - """ - Handle the crawling process for a single URL. - """ - documents = [] - start_time = asyncio.get_event_loop().time() - - async with session.post( - crawl_url, headers=headers, json=request_data - ) as response: - data = await self.handle_response(response) - job_id = data.get("jobId") - msg.info(f"Creating Firecrawl Job {job_id}") - - if job_id: - documents = await self.poll_job_status( - session, crawl_url, headers, job_id, start_time - ) - - return documents - - async def poll_job_status( - self, - session: aiohttp.ClientSession, - crawl_url: str, - headers: dict, - job_id: str, - start_time: float, - ) -> List[Tuple[str, str, str]]: - """ - Poll the job status and retrieve results when completed. - """ - max_retries = 60 - wait_time = 10 - documents = [] - - for attempt in range(max_retries): - elapsed_time = round(asyncio.get_event_loop().time() - start_time, 2) - msg.info( - f"Checking Firecrawl Job Status for {job_id} (Try: {attempt + 1}) ({elapsed_time}s)" - ) - - async with session.get( - f"{crawl_url}/status/{job_id}", headers=headers - ) as response: - data = await self.handle_response(response) - status = data.get("status") - msg.info(f"Firecrawl Job Status: {status}") - - files = data.get("data", []) - if files: - msg.info(f"{len(files)} Files scraped") - - if status == "completed": - msg.good("Firecrawl Job successful") - documents.extend( - ( - file["metadata"]["title"], - file["markdown"], - file["metadata"]["sourceURL"], - ) - for file in files - ) - break - - if attempt < max_retries - 1: - await asyncio.sleep(wait_time) - - return documents diff --git a/goldenverba/components/reader/GitReader.py b/goldenverba/components/reader/GitReader.py index f2c76b150f..a231763c7a 100644 --- a/goldenverba/components/reader/GitReader.py +++ b/goldenverba/components/reader/GitReader.py @@ -117,6 +117,7 @@ async def load(self, config: dict, fileConfig: FileConfig) -> list[Document]: file_size=size, status=fileConfig.status, status_report=fileConfig.status_report, + metadata=fileConfig.metadata, ) document = await reader.load(config, new_file_config) documents.append(document[0]) @@ -135,7 +136,8 @@ async def fetch_docs_github( self, url: str, folder: str, token: str, reader: Reader ) -> list[str]: headers = self.get_headers(token, "GitHub") - async with aiohttp.ClientSession() as session: + timeout = aiohttp.ClientTimeout(total=30) + async with aiohttp.ClientSession(timeout=timeout) as session: async with session.get(url, headers=headers) as response: response.raise_for_status() data = await response.json() @@ -148,7 +150,8 @@ async def fetch_docs_github( async def fetch_docs_gitlab(self, url: str, token: str, reader: Reader) -> list: headers = self.get_headers(token, "GitLab") - async with aiohttp.ClientSession() as session: + timeout = aiohttp.ClientTimeout(total=30) + async with aiohttp.ClientSession(timeout=timeout) as session: async with session.get(url, headers=headers) as response: response.raise_for_status() data = await response.json() @@ -166,7 +169,8 @@ async def download_file_github( f"https://api.github.com/repos/{owner}/{name}/contents/{path}?ref={branch}" ) headers = self.get_headers(token, "GitHub") - async with aiohttp.ClientSession() as session: + timeout = aiohttp.ClientTimeout(total=30) + async with aiohttp.ClientSession(timeout=timeout) as session: async with session.get(url, headers=headers) as response: response.raise_for_status() data = await response.json() @@ -183,7 +187,8 @@ async def download_file_gitlab( url = f"https://gitlab.com/api/v4/projects/{project_id}/repository/files/{urllib.parse.quote(file_path, safe='')}/raw?ref={branch}" headers = {"PRIVATE-TOKEN": token} - async with aiohttp.ClientSession() as session: + timeout = aiohttp.ClientTimeout(total=30) + async with aiohttp.ClientSession(timeout=timeout) as session: async with session.get(url, headers=headers) as response: if response.status == 200: content = await response.read() diff --git a/goldenverba/components/reader/HTMLReader.py b/goldenverba/components/reader/HTMLReader.py index c29423303a..4f383a39f2 100644 --- a/goldenverba/components/reader/HTMLReader.py +++ b/goldenverba/components/reader/HTMLReader.py @@ -69,7 +69,8 @@ async def load(self, config: dict, fileConfig: FileConfig) -> list[Document]: documents = [] processed_urls = set() - async with aiohttp.ClientSession() as session: + timeout = aiohttp.ClientTimeout(total=60) + async with aiohttp.ClientSession(timeout=timeout) as session: for url in urls: try: await self.process_url( @@ -80,6 +81,7 @@ async def load(self, config: dict, fileConfig: FileConfig) -> list[Document]: 0, session, reader, + config, fileConfig, documents, processed_urls, @@ -98,6 +100,7 @@ async def process_url( current_depth: int, session: aiohttp.ClientSession, reader: BasicReader, + config: dict, fileConfig: FileConfig, documents: List[Document], processed_urls: set, @@ -126,7 +129,7 @@ async def process_url( status_report=fileConfig.status_report, metadata=fileConfig.metadata, ) - document = await reader.load(self.config, new_file_config) + document = await reader.load(config, new_file_config) documents.extend(document) if recursive and current_depth < max_depth: @@ -140,6 +143,7 @@ async def process_url( current_depth + 1, session, reader, + config, fileConfig, documents, processed_urls, diff --git a/goldenverba/components/reader/UnstructuredAPI.py b/goldenverba/components/reader/UnstructuredAPI.py index 57c8648637..fc2e166a09 100644 --- a/goldenverba/components/reader/UnstructuredAPI.py +++ b/goldenverba/components/reader/UnstructuredAPI.py @@ -2,7 +2,6 @@ import io import os -import requests from wasabi import msg import aiohttp @@ -89,7 +88,8 @@ async def load( ) try: - async with aiohttp.ClientSession() as session: + timeout = aiohttp.ClientTimeout(total=120) + async with aiohttp.ClientSession(timeout=timeout) as session: async with session.post( api_url, headers=headers, data=file_data ) as response: @@ -105,7 +105,7 @@ async def load( return [create_document(file_content, fileConfig)] - except requests.RequestException as e: + except aiohttp.ClientError as e: raise Exception( f"Unstructured API request failed for {fileConfig.filename}: {str(e)}" ) diff --git a/goldenverba/components/reader/UpstageDocumentParse.py b/goldenverba/components/reader/UpstageDocumentParse.py deleted file mode 100644 index d83ad54797..0000000000 --- a/goldenverba/components/reader/UpstageDocumentParse.py +++ /dev/null @@ -1,96 +0,0 @@ -import base64 -import io -import os - -import requests -from wasabi import msg -import aiohttp - -from goldenverba.components.document import Document, create_document -from goldenverba.components.interfaces import Reader -from goldenverba.server.types import FileConfig -from goldenverba.components.util import get_environment -from goldenverba.components.types import InputConfig - - -class UpstageDocumentParseReader(Reader): - """ - Upstage Document Parse API Reader for converting documents to structured HTML format. - """ - - def __init__(self): - super().__init__() - self.requires_env = ["UPSTAGE_API_KEY"] - self.name = "Upstage Parser" - self.description = "Uses the Upstage Document Parse API to convert documents into structured HTML format" - - if os.getenv("UPSTAGE_API_KEY") is None: - self.config["API Key"] = InputConfig( - type="password", - value="", - description="Set your Upstage API Key here or set it as an environment variable `UPSTAGE_API_KEY`", - values=[], - ) - - if os.getenv("UPSTAGE_API_URL") is None: - self.config["API URL"] = InputConfig( - type="text", - value="https://api.upstage.ai/v1/document-ai/document-parse", - description="Set the base URL to the Upstage API", - values=[], - ) - - async def load( - self, config: dict[str, InputConfig], fileConfig: FileConfig - ) -> list[Document]: - """ - Load and process a file using the Upstage Document Parse API. - """ - # Get API credentials - token = get_environment( - config, - "API Key", - "UPSTAGE_API_KEY", - "No Upstage API Key detected", - ) - api_url = get_environment( - config, "API URL", "UPSTAGE_API_URL", "No Upstage API URL detected" - ) - - headers = { - "Authorization": f"Bearer {token}", - } - - msg.info(f"Loading {fileConfig.filename}") - - file_data = aiohttp.FormData() - file_bytes = io.BytesIO(base64.b64decode(fileConfig.content)) - file_data.add_field( - "document", - file_bytes, - filename=f"{fileConfig.filename}.{fileConfig.extension}", - ) - - try: - async with aiohttp.ClientSession() as session: - async with session.post( - api_url, headers=headers, data=file_data - ) as response: - response.raise_for_status() - json_response = await response.json() - - if "content" not in json_response: - raise ValueError(f"API error: Invalid response format") - - # Extract text content from HTML - html_content = json_response["content"]["html"] - # You might want to add HTML to text conversion here - # For now, we'll use the HTML content directly - return [create_document(html_content, fileConfig)] - - except aiohttp.ClientError as e: - raise Exception( - f"Upstage API request failed for {fileConfig.filename}: {str(e)}" - ) - except Exception as e: - raise Exception(f"Failed to process {fileConfig.filename}: {str(e)}") diff --git a/goldenverba/components/reader/WhisperReader.py b/goldenverba/components/reader/WhisperReader.py new file mode 100644 index 0000000000..52ffe0aab4 --- /dev/null +++ b/goldenverba/components/reader/WhisperReader.py @@ -0,0 +1,126 @@ +import asyncio +import base64 +import os +import tempfile + +from wasabi import msg + +from goldenverba.components.document import Document, create_document +from goldenverba.components.interfaces import Reader +from goldenverba.server.types import FileConfig +from goldenverba.components.types import InputConfig + +try: + from faster_whisper import WhisperModel +except ImportError: + WhisperModel = None + + +class WhisperReader(Reader): + """ + Local Whisper reader for importing audio and video files using faster-whisper. + Runs entirely locally — no API key or external service required. + """ + + def __init__(self): + super().__init__() + self.requires_library = ["faster_whisper"] + self.extension = [ + ".3ga", + ".8svx", + ".aac", + ".ac3", + ".aif", + ".aiff", + ".alac", + ".amr", + ".ape", + ".au", + ".dss", + ".flac", + ".flv", + ".m2ts", + ".m4a", + ".m4b", + ".m4p", + ".m4r", + ".m4v", + ".mov", + ".mp2", + ".mp3", + ".mp4", + ".mpga", + ".mts", + ".mxf", + ".ogg", + ".oga", + ".mogg", + ".opus", + ".qcp", + ".ts", + ".tta", + ".voc", + ".wav", + ".webm", + ".wma", + ".wv", + ] + self.name = "Whisper" + self.description = "Transcribes audio and video files locally using faster-whisper. No API key required." + self.config = { + "Model Size": InputConfig( + type="dropdown", + value="base", + description="Whisper model size — larger models are more accurate but slower and use more memory", + values=["tiny", "base", "small", "medium", "large-v3"], + ), + "Device": InputConfig( + type="dropdown", + value="cpu", + description="Compute device for inference", + values=["cpu", "cuda", "auto"], + ), + } + + async def load( + self, config: dict[str, InputConfig], fileConfig: FileConfig + ) -> list[Document]: + """ + Transcribe an audio/video file using faster-whisper running locally. + """ + if WhisperModel is None: + raise ImportError( + "faster-whisper is required for audio transcription. " + "Install it with: pip install faster-whisper" + ) + + model_size = config["Model Size"].value + device = config["Device"].value + + msg.info(f"Transcribing {fileConfig.filename} with Whisper ({model_size})") + + file_bytes = base64.b64decode(fileConfig.content) + + # faster-whisper needs a file path, not a file-like object + suffix = f".{fileConfig.extension}" if fileConfig.extension else ".wav" + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp: + tmp.write(file_bytes) + tmp_path = tmp.name + + try: + segments, _info = await asyncio.to_thread( + self._transcribe, model_size, device, tmp_path + ) + text = " ".join(segment.text.strip() for segment in segments) + if not text: + raise Exception(f"Whisper returned empty transcript for {fileConfig.filename}") + return [create_document(text, fileConfig)] + finally: + os.unlink(tmp_path) + + def _transcribe(self, model_size: str, device: str, path: str): + """Synchronous transcription — called via asyncio.to_thread.""" + model = WhisperModel(model_size, device=device, compute_type="int8") + segments, info = model.transcribe(path, beam_size=5) + # Consume the generator inside the thread so it doesn't escape to async context + return list(segments), info diff --git a/goldenverba/components/reader/reader_manager.py b/goldenverba/components/reader/reader_manager.py new file mode 100644 index 0000000000..8b0921e828 --- /dev/null +++ b/goldenverba/components/reader/reader_manager.py @@ -0,0 +1,78 @@ +""" +reader_manager.py +================= +Reader component registry and ReaderManager. + +To add a new Reader: + 1. Implement it in this directory (goldenverba/components/reader/) + 2. Import it below and add an instance to the `readers` list +""" + +import asyncio + +from goldenverba.components.document import Document +from goldenverba.components.interfaces import Reader +from goldenverba.server.helpers import LoggerManager +from goldenverba.server.types import FileConfig, FileStatus + +from goldenverba.components.reader.BasicReader import BasicReader +from goldenverba.components.reader.GitReader import GitReader +from goldenverba.components.reader.UnstructuredAPI import UnstructuredReader +from goldenverba.components.reader.HTMLReader import HTMLReader +from goldenverba.components.reader.WhisperReader import WhisperReader + +# All available readers — add new instances here +readers = [ + BasicReader(), + HTMLReader(), + GitReader(), + UnstructuredReader(), + WhisperReader(), +] + + +class ReaderManager: + """Dispatches load() calls to the correct Reader implementation.""" + + def __init__(self): + self.readers: dict[str, Reader] = {reader.name: reader for reader in readers} + + async def load( + self, reader: str, fileConfig: FileConfig, logger: LoggerManager + ) -> list[Document]: + try: + loop = asyncio.get_running_loop() + start_time = loop.time() + if reader in self.readers: + config = fileConfig.rag_config["Reader"].components[reader].config + documents: list[Document] = await self.readers[reader].load( + config, fileConfig + ) + for document in documents: + document.meta["Reader"] = ( + fileConfig.rag_config["Reader"].components[reader].model_dump() + ) + elapsed_time = round(loop.time() - start_time, 2) + if len(documents) == 1: + await logger.send_report( + fileConfig.fileID, + FileStatus.LOADING, + f"Loaded {fileConfig.filename}", + took=elapsed_time, + ) + else: + await logger.send_report( + fileConfig.fileID, + FileStatus.LOADING, + f"Loaded {fileConfig.filename} with {len(documents)} documents", + took=elapsed_time, + ) + await logger.send_report( + fileConfig.fileID, FileStatus.CHUNKING, "", took=0 + ) + return documents + else: + raise Exception(f"{reader} Reader not found") + + except Exception as e: + raise Exception(f"Reader {reader} failed with: {str(e)}") diff --git a/goldenverba/components/retriever/retriever_manager.py b/goldenverba/components/retriever/retriever_manager.py new file mode 100644 index 0000000000..d9bf03f4f5 --- /dev/null +++ b/goldenverba/components/retriever/retriever_manager.py @@ -0,0 +1,62 @@ +""" +retriever_manager.py +==================== +Retriever component registry and RetrieverManager. + +To add a new Retriever: + 1. Implement it in this directory (goldenverba/components/retriever/) + 2. Import it below and add an instance to the `retrievers` list +""" + +from goldenverba.components.interfaces import Retriever +from goldenverba.components.weaviate_manager import WeaviateManager +from goldenverba.components.retriever.WindowRetriever import WindowRetriever + +# All available retrievers — add new instances here +retrievers = [WindowRetriever()] + + +class RetrieverManager: + """Dispatches retrieve() calls to the correct Retriever implementation.""" + + def __init__(self): + self.retrievers: dict[str, Retriever] = { + retriever.name: retriever for retriever in retrievers + } + + async def retrieve( + self, + client, + retriever: str, + query: str, + vector: list[float], + rag_config: dict, + weaviate_manager: WeaviateManager, + labels: list[str], + document_uuids: list[str], + ): + try: + if retriever not in self.retrievers: + raise Exception(f"Retriever {retriever} not found") + + embedder_model = ( + rag_config["Embedder"] + .components[rag_config["Embedder"].selected] + .config["Model"] + .value + ) + config = rag_config["Retriever"].components[retriever].config + documents, context = await self.retrievers[retriever].retrieve( + client, + query, + vector, + config, + weaviate_manager, + embedder_model, + labels, + document_uuids, + ) + return (documents, context) + + except Exception as e: + raise e diff --git a/goldenverba/components/util.py b/goldenverba/components/util.py index f376e25051..e9a2bf9f12 100644 --- a/goldenverba/components/util.py +++ b/goldenverba/components/util.py @@ -33,13 +33,9 @@ def transform_data(X, components): # Function to perform PCA def pca(X, k): - print(X[:10]) X_standardized = standardize_data(X) - print(X_standardized[:10]) covariance_matrix = compute_covariance_matrix(X_standardized) - print(covariance_matrix) eigenvalues, eigenvectors = eigen_decomposition(covariance_matrix) - print(eigenvalues, eigenvectors) sorted_eigenvalues, sorted_eigenvectors = sort_eigenvalues_eigenvectors(eigenvalues, eigenvectors) top_k_components = select_top_k_components(sorted_eigenvectors, k) X_pca = transform_data(X_standardized, top_k_components) diff --git a/goldenverba/components/verba_manager.py b/goldenverba/components/verba_manager.py new file mode 100644 index 0000000000..1e1af2c381 --- /dev/null +++ b/goldenverba/components/verba_manager.py @@ -0,0 +1,660 @@ +""" +verba_manager.py +================ +Core pipeline orchestrator for Verba. + + VerbaManager — owns one instance of every component manager (Reader, Chunker, + Embedder, Retriever, Generator, Weaviate) and wires them together + for document import, retrieval, and generation. + + One VerbaManager lives inside ClientManager (client_manager.py), + which handles connection pooling. api.py also keeps a module-level + VerbaManager singleton for config-only calls that don't need a + live Weaviate connection. + +Adding a new pipeline component (e.g. a new Generator)? + 1. Implement it in goldenverba/components/generation/ + 2. Register it in GeneratorManager (components/generation/generator_manager.py) + 3. Nothing else needs changing here. +""" + +import os +import importlib +import math +import json + +from dotenv import load_dotenv +from wasabi import msg +import asyncio + +from copy import deepcopy +from goldenverba.server.helpers import LoggerManager + +from goldenverba.components.document import Document +from goldenverba.server.types import ( + FileConfig, + FileStatus, + ChunkScore, + Credentials, +) + +from goldenverba.components.reader.reader_manager import ReaderManager +from goldenverba.components.chunking.chunker_manager import ChunkerManager +from goldenverba.components.embedding.embedding_manager import EmbeddingManager +from goldenverba.components.retriever.retriever_manager import RetrieverManager +from goldenverba.components.generation.generator_manager import GeneratorManager +from goldenverba.components.weaviate_manager import WeaviateManager + +load_dotenv() + + +class VerbaManager: + """ + Orchestrates the full Verba pipeline. + + Holds one instance of each component manager and exposes high-level + async methods used by the FastAPI layer (api.py) and ClientManager. + Does not manage Weaviate connections directly — that is ClientManager's job. + """ + + def __init__(self) -> None: + # One instance of each component manager; each manager holds the full + # registry of available implementations (e.g. all Generator subclasses). + self.reader_manager = ReaderManager() + self.chunker_manager = ChunkerManager() + self.embedder_manager = EmbeddingManager() + self.retriever_manager = RetrieverManager() + self.generator_manager = GeneratorManager() + self.weaviate_manager = WeaviateManager() + + # Fixed UUIDs for the three config documents stored in Weaviate. + # These never change so that configs survive restarts. + self.rag_config_uuid = "e0adcc12-9bad-4588-8a1e-bab0af6ed485" + self.theme_config_uuid = "baab38a7-cb51-4108-acd8-6edeca222820" + self.user_config_uuid = "f53f7738-08be-4d5a-b003-13eb4bf03ac7" + + # Populated at startup; passed to components so they can mark themselves + # available/unavailable in the UI without attempting live calls. + self.environment_variables: dict[str, bool] = {} + self.installed_libraries: dict[str, bool] = {} + + self.verify_installed_libraries() + self.verify_variables() + + # ------------------------------------------------------------------------- + # Connection + # ------------------------------------------------------------------------- + + async def connect(self, credentials: Credentials, port: str = "8080"): + """Open a Weaviate client and ensure the config collection exists.""" + loop = asyncio.get_running_loop() + start_time = loop.time() + client = await self.weaviate_manager.connect( + credentials.deployment, credentials.url, credentials.key, port + ) + if client: + initialized = await self.weaviate_manager.verify_collection( + client, self.weaviate_manager.config_collection_name + ) + if initialized: + msg.info(f"Connection time: {loop.time() - start_time:.2f} seconds") + return client + raise Exception( + "Connected to Weaviate but failed to verify configuration collection" + ) + raise Exception("Weaviate client could not be created") + + async def disconnect(self, client): + """Close a Weaviate client connection.""" + loop = asyncio.get_running_loop() + start_time = loop.time() + result = await self.weaviate_manager.disconnect(client) + msg.info(f"Disconnection time: {loop.time() - start_time:.2f} seconds") + return result + + async def get_deployments(self): + """Return Weaviate connection env vars so the frontend can pre-fill them.""" + return { + "WEAVIATE_URL_VERBA": os.getenv("WEAVIATE_URL_VERBA") or "", + "WEAVIATE_API_KEY_VERBA": os.getenv("WEAVIATE_API_KEY_VERBA") or "", + } + + # ------------------------------------------------------------------------- + # Import pipeline + # ------------------------------------------------------------------------- + + async def import_document( + self, client, fileConfig: FileConfig, logger: LoggerManager = None + ): + """ + Entry point for ingesting one file/URL. + + Flow: duplicate check → Reader.load() → process_single_document() per doc. + All per-document tasks run concurrently via asyncio.gather(). + Progress is streamed back to the caller via `logger` (a WebSocket wrapper). + """ + if logger is None: + logger = LoggerManager() + try: + loop = asyncio.get_running_loop() + start_time = loop.time() + + # Check for an existing document with the same name. + duplicate_uuid = await self.weaviate_manager.exist_document_name( + client, fileConfig.filename + ) + if duplicate_uuid is not None and not fileConfig.overwrite: + raise Exception(f"{fileConfig.filename} already exists in Verba") + elif duplicate_uuid is not None and fileConfig.overwrite: + await self.weaviate_manager.delete_document(client, duplicate_uuid) + await logger.send_report( + fileConfig.fileID, + status=FileStatus.STARTING, + message=f"Overwriting {fileConfig.filename}", + took=0, + ) + else: + await logger.send_report( + fileConfig.fileID, + status=FileStatus.STARTING, + message="Starting Import", + took=0, + ) + + # Reader turns the raw file/URL into one or more Document objects. + documents = await self.reader_manager.load( + fileConfig.rag_config["Reader"].selected, fileConfig, logger + ) + + # Process all documents concurrently; collect exceptions instead of failing fast. + tasks = [ + self.process_single_document(client, doc, fileConfig, logger) + for doc in documents + ] + results = await asyncio.gather(*tasks, return_exceptions=True) + successful_tasks = sum( + 1 for result in results if not isinstance(result, Exception) + ) + + # Report outcome: multi-doc (e.g. URL with multiple pages), single doc, or failure. + if successful_tasks > 1: + await logger.send_report( + fileConfig.fileID, + status=FileStatus.INGESTING, + message=f"Imported {fileConfig.filename} and it's {successful_tasks} documents into Weaviate", + took=round(loop.time() - start_time, 2), + ) + elif successful_tasks == 1: + await logger.send_report( + fileConfig.fileID, + status=FileStatus.INGESTING, + message=f"Imported {fileConfig.filename} and {len(documents[0].chunks)} chunks into Weaviate", + took=round(loop.time() - start_time, 2), + ) + elif ( + successful_tasks == 0 + and len(results) == 1 + and isinstance(results[0], Exception) + ): + msg.fail( + f"No documents imported {successful_tasks} of {len(results)} succesful tasks" + ) + raise results[0] + else: + raise Exception( + f"No documents imported {successful_tasks} of {len(results)} succesful tasks" + ) + + await logger.send_report( + fileConfig.fileID, + status=FileStatus.DONE, + message=f"Import for {fileConfig.filename} completed successfully", + took=round(loop.time() - start_time, 2), + ) + + except Exception as e: + await logger.send_report( + fileConfig.fileID, + status=FileStatus.ERROR, + message=f"Import for {fileConfig.filename} failed: {str(e)}", + took=0, + ) + return + + async def process_single_document( + self, + client, + document: Document, + fileConfig: FileConfig, + logger: LoggerManager, + ): + """ + Chunk → embed → store one Document in Weaviate. + + For URL imports a single FileConfig can expand into multiple Documents + (e.g. one per page). Each gets its own derived fileID so the frontend + can track them independently. + """ + loop = asyncio.get_running_loop() + start_time = loop.time() + + # URL imports: each extracted document gets its own config and logger entry. + if fileConfig.isURL: + currentFileConfig = deepcopy(fileConfig) + currentFileConfig.fileID = fileConfig.fileID + document.title + currentFileConfig.isURL = False + currentFileConfig.filename = document.title + await logger.create_new_document( + fileConfig.fileID + document.title, + document.title, + fileConfig.fileID, + ) + else: + currentFileConfig = fileConfig + + try: + duplicate_uuid = await self.weaviate_manager.exist_document_name( + client, document.title + ) + if duplicate_uuid is not None and not currentFileConfig.overwrite: + raise Exception(f"{document.title} already exists in Verba") + elif duplicate_uuid is not None and currentFileConfig.overwrite: + await self.weaviate_manager.delete_document(client, duplicate_uuid) + + # Chunker splits the document into overlapping text windows. + chunked_documents = await self.chunker_manager.chunk( + currentFileConfig.rag_config["Chunker"].selected, + currentFileConfig, + [document], + self.embedder_manager.embedders[ + currentFileConfig.rag_config["Embedder"].selected + ], + logger, + ) + + # Embedder turns each chunk into a vector; may batch-call an external API. + vectorized_documents = await self.embedder_manager.vectorize( + currentFileConfig.rag_config["Embedder"].selected, + currentFileConfig, + chunked_documents, + logger, + ) + + # Write each vectorized document (and its chunks) to Weaviate. + for document in vectorized_documents: + await self.weaviate_manager.import_document( + client, + document, + currentFileConfig.rag_config["Embedder"] + .components[fileConfig.rag_config["Embedder"].selected] + .config["Model"] + .value, + ) + + await logger.send_report( + currentFileConfig.fileID, + status=FileStatus.INGESTING, + message=f"Imported {currentFileConfig.filename} into Weaviate", + took=round(loop.time() - start_time, 2), + ) + + await logger.send_report( + currentFileConfig.fileID, + status=FileStatus.DONE, + message=f"Import for {currentFileConfig.filename} completed successfully", + took=round(loop.time() - start_time, 2), + ) + except Exception as e: + await logger.send_report( + currentFileConfig.fileID, + status=FileStatus.ERROR, + message=f"Import for {fileConfig.filename} failed: {str(e)}", + took=round(loop.time() - start_time, 2), + ) + raise Exception(f"Import for {fileConfig.filename} failed: {str(e)}") + + # ------------------------------------------------------------------------- + # Configuration + # ------------------------------------------------------------------------- + + def _build_category_config(self, components: dict) -> dict: + """ + Build the config payload for one pipeline category (e.g. "Reader"). + + Returns {"components": {name: meta, ...}, "selected": first_component_name}. + `get_meta()` on each component includes env/library availability so the + frontend knows which components are usable without an extra round-trip. + """ + return { + "components": { + k: v.get_meta(self.environment_variables, self.installed_libraries) + for k, v in components.items() + }, + "selected": next(iter(components.values())).name, + } + + def create_config(self) -> dict: + """ + Build a fresh RAG config from the currently registered components. + + Called on every load_rag_config() to compare against the stored config + and detect schema drift (added/removed components or config keys). + """ + return { + "Reader": self._build_category_config(self.reader_manager.readers), + "Chunker": self._build_category_config(self.chunker_manager.chunkers), + "Embedder": self._build_category_config(self.embedder_manager.embedders), + "Retriever": self._build_category_config(self.retriever_manager.retrievers), + "Generator": self._build_category_config(self.generator_manager.generators), + } + + def create_user_config(self) -> dict: + """Default user config returned when no stored config exists yet.""" + return {"getting_started": False} + + # Thin pass-throughs — config storage lives in Weaviate, not in memory. + async def set_theme_config(self, client, config: dict): + await self.weaviate_manager.set_config(client, self.theme_config_uuid, config) + + async def set_rag_config(self, client, config: dict): + await self.weaviate_manager.set_config(client, self.rag_config_uuid, config) + + async def set_user_config(self, client, config: dict): + await self.weaviate_manager.set_config(client, self.user_config_uuid, config) + + async def load_rag_config(self, client): + """ + Return a valid RAG config, preferring the stored one. + + If the stored config is missing or fails verify_config() (schema drift), + fall back to a freshly generated config and persist it. + """ + loaded_config = await self.weaviate_manager.get_config( + client, self.rag_config_uuid + ) + new_config = self.create_config() + if loaded_config is not None: + if self.verify_config(loaded_config, new_config): + msg.info("Using Existing RAG Configuration") + return loaded_config + else: + msg.info("Using New RAG Configuration") + await self.set_rag_config(client, new_config) + return new_config + else: + msg.info("Using New RAG Configuration") + return new_config + + async def load_theme_config(self, client): + """Return (theme, themes) from Weaviate, or (None, None) if not set.""" + loaded_config = await self.weaviate_manager.get_config( + client, self.theme_config_uuid + ) + if loaded_config is None: + return None, None + return loaded_config["theme"], loaded_config["themes"] + + async def load_user_config(self, client): + """Return the stored user config, or a fresh default if none exists.""" + loaded_config = await self.weaviate_manager.get_config( + client, self.user_config_uuid + ) + if loaded_config is None: + return self.create_user_config() + return loaded_config + + @staticmethod + def _keys_match(a: dict, b: dict, label: str) -> bool: + """Return True if both dicts have identical key sets; log and return False otherwise.""" + if set(a.keys()) == set(b.keys()): + return True + msg.fail(f"Config Validation Failed, {label}: {set(a.keys())} != {set(b.keys())}") + return False + + def verify_config(self, a: dict, b: dict) -> bool: + """ + Compare stored config `a` against authoritative config `b` (4 levels deep). + + Walks categories → components → config keys → setting fields. + Returns False on the first mismatch so the caller knows to regenerate. + In Demo mode, always returns True to avoid overwriting a shared config. + """ + try: + if os.getenv("VERBA_PRODUCTION") == "Demo": + return True + + if not self._keys_match(a, b, "category mismatch"): + return False + + for category_key in b: + a_components = a[category_key]["components"] + b_components = b[category_key]["components"] + if not self._keys_match(a_components, b_components, f"{category_key} component mismatch"): + return False + + for component_key in b_components: + a_config = a_components[component_key]["config"] + b_config = b_components[component_key]["config"] + if not self._keys_match(a_config, b_config, f"{component_key} config key mismatch"): + return False + + for config_key in b_config: + a_s, b_s = a_config[config_key], b_config[config_key] + if a_s["description"] != b_s["description"]: + msg.fail(f"Config Validation Failed, description mismatch: {a_s['description']} != {b_s['description']}") + return False + if sorted(a_s["values"]) != sorted(b_s["values"]): + msg.fail(f"Config Validation Failed, values mismatch: {a_s['values']} != {b_s['values']}") + return False + + return True + + except Exception as e: + msg.fail(f"Config Validation failed: {str(e)}") + return False + + async def reset_rag_config(self, client): + msg.info("Resetting RAG Configuration") + await self.weaviate_manager.reset_config(client, self.rag_config_uuid) + + async def reset_theme_config(self, client): + msg.info("Resetting Theme Configuration") + await self.weaviate_manager.reset_config(client, self.theme_config_uuid) + + async def reset_user_config(self, client): + msg.info("Resetting User Configuration") + await self.weaviate_manager.reset_config(client, self.user_config_uuid) + + # ------------------------------------------------------------------------- + # Environment and library introspection + # ------------------------------------------------------------------------- + + def _collect_from_managers(self, attr: str) -> set[str]: + """ + Union of `attr` (e.g. 'requires_library') across every registered component. + + Used to build the availability maps shown in the status page without + duplicating component iteration logic in each verify_* method. + """ + managers = [ + (self.reader_manager, "readers"), + (self.chunker_manager, "chunkers"), + (self.embedder_manager, "embedders"), + (self.retriever_manager, "retrievers"), + (self.generator_manager, "generators"), + ] + return { + item + for mgr, collection_attr in managers + for component in getattr(mgr, collection_attr).values() + for item in getattr(component, attr) + } + + def verify_installed_libraries(self) -> None: + """ + Attempt to import every library declared by any component. + Populates self.installed_libraries {lib_name: bool}. + """ + for lib in self._collect_from_managers("requires_library"): + try: + importlib.import_module(lib) + self.installed_libraries[lib] = True + except Exception: + self.installed_libraries[lib] = False + + def verify_variables(self) -> None: + """ + Check which env vars declared by any component are actually set. + Populates self.environment_variables {var_name: bool}. + """ + for env in self._collect_from_managers("requires_env"): + self.environment_variables[env] = os.environ.get(env) is not None + + # ------------------------------------------------------------------------- + # Document content retrieval + # ------------------------------------------------------------------------- + + async def get_content( + self, + client, + uuid: str, + page: int, + chunkScores: list[ChunkScore], + ): + """ + Return paginated document content for the Document Explorer. + + Two modes: + chunkScores present — RAG mode. Shows the matched chunk plus up to 5 + surrounding chunks for context. Three Weaviate fetches run in parallel. + chunkScores empty — Browse mode. Returns one page of sequential chunks + and the total page count (chunk count + chunk count run in parallel). + + Returns: (content_pieces, total_batches) + content_pieces: list of {"content", "chunk_id", "score", "type"} dicts + total_batches: total number of pages/scores available + """ + chunks_per_page = 10 + content_pieces = [] + total_batches = 0 + + if len(chunkScores) > 0: + # RAG mode: show the matched chunk with surrounding context window. + if page > len(chunkScores): + page = 0 + + total_batches = len(chunkScores) + score = chunkScores[page] + half = chunks_per_page // 2 + + before_ids = list(range(max(0, score.chunk_id - half), score.chunk_id)) + after_ids = list(range(score.chunk_id + 1, score.chunk_id + half)) + + async def _empty(): + return [] + + # Fetch target + before/after context in a single round-trip. + chunk, chunks_before, chunks_after = await asyncio.gather( + self.weaviate_manager.get_chunk(client, score.uuid, score.embedder), + self.weaviate_manager.get_chunk_by_ids(client, score.embedder, uuid, before_ids) + if before_ids + else _empty(), + self.weaviate_manager.get_chunk_by_ids(client, score.embedder, uuid, after_ids) + if after_ids + else _empty(), + ) + + before_content = "".join( + c.properties["content_without_overlap"] for c in (chunks_before or []) + ) + after_content = "".join( + c.properties["content_without_overlap"] for c in (chunks_after or []) + ) + + content_pieces.append({"content": before_content, "chunk_id": 0, "score": 0, "type": "text"}) + content_pieces.append({"content": chunk["content_without_overlap"] if chunk else "", "chunk_id": score.chunk_id, "score": score.score, "type": "extract"}) + content_pieces.append({"content": after_content, "chunk_id": 0, "score": 0, "type": "text"}) + + else: + # Browse mode: return one sequential page of chunks. + document = await self.weaviate_manager.get_document( + client, uuid, properties=["meta"] + ) + if not document or not document.get("meta"): + return (content_pieces, total_batches) + + config = json.loads(document["meta"]) + embedder = config["Embedder"]["config"]["Model"]["value"] + request_chunk_ids = list(range(chunks_per_page * page, chunks_per_page * (page + 1))) + + # Fetch the page content and total count in parallel. + chunks, total_chunks = await asyncio.gather( + self.weaviate_manager.get_chunk_by_ids(client, embedder, uuid, request_chunk_ids), + self.weaviate_manager.get_chunk_count(client, embedder, uuid), + ) + total_batches = int(math.ceil(total_chunks / chunks_per_page)) + content = "".join(chunk.properties["content_without_overlap"] for chunk in chunks) + content_pieces.append({"content": content, "chunk_id": 0, "score": 0, "type": "text"}) + + return (content_pieces, total_batches) + + # ------------------------------------------------------------------------- + # RAG pipeline + # ------------------------------------------------------------------------- + + async def retrieve_chunks( + self, + client, + query: str, + rag_config: dict, + labels: list[str] | None = None, + document_uuids: list[str] | None = None, + ): + """ + Embed the query and retrieve relevant chunks via the selected Retriever. + + Also writes the query as an autocomplete suggestion (skipped for very + short queries to avoid polluting suggestions with partial keystrokes). + + Returns: (documents, context_string) + """ + labels = labels or [] + document_uuids = document_uuids or [] + + retriever = rag_config["Retriever"].selected + embedder = rag_config["Embedder"].selected + + # Only persist as a suggestion if the query is meaningful. + if query and len(query.strip()) >= 3: + await self.weaviate_manager.add_suggestion(client, query) + + vector = await self.embedder_manager.vectorize_query(embedder, query, rag_config) + documents, context = await self.retriever_manager.retrieve( + client, retriever, query, vector, rag_config, + self.weaviate_manager, labels, document_uuids, + ) + + return (documents, context) + + async def generate_stream_answer( + self, + rag_config: dict, + query: str, + context: str, + conversation: list[dict], + ): + """ + Async generator that streams token chunks from the selected Generator. + + Accumulates tokens into full_text_parts and attaches full_text to the + final "stop" chunk so the caller gets the complete answer in one place. + Yields each result dict directly to the WebSocket handler in api.py. + """ + full_text_parts: list[str] = [] + async for result in self.generator_manager.generate_stream( + rag_config, query, context, conversation + ): + full_text_parts.append(result["message"]) + if result.get("finish_reason") == "stop": + result["full_text"] = "".join(full_text_parts) + yield result diff --git a/goldenverba/components/managers.py b/goldenverba/components/weaviate_manager.py similarity index 50% rename from goldenverba/components/managers.py rename to goldenverba/components/weaviate_manager.py index e6454d0819..850e64372f 100644 --- a/goldenverba/components/managers.py +++ b/goldenverba/components/weaviate_manager.py @@ -1,12 +1,36 @@ +""" +weaviate_manager.py +=================== +Low-level Weaviate client wrapper for Verba. + +Owns all direct interactions with the Weaviate async client: connections, +collection management, document/chunk CRUD, vector queries, suggestions, +and configuration storage. + +Collection layout +----------------- +VERBA_DOCUMENTS — one record per imported file (title, meta, labels) +VERBA_CONFIGURATION — three records: RAG config, theme config, user config +VERBA_SUGGESTIONS — autocomplete query history +VERBA_Embedding_<model> — one collection per embedder model (chunks + vectors) +VERBA_Cache_<model> — semantic-cache entries per embedder model + +Deployment modes +---------------- +"Weaviate" — Weaviate Cloud (URL + API key required) +"Docker" — local Docker Compose instance (default host: localhost or WEAVIATE_HOST env var) +"Custom" — any reachable Weaviate instance with optional auth +""" + from wasabi import msg import weaviate from weaviate.client import WeaviateAsyncClient -from weaviate.auth import AuthApiKey from weaviate.classes.query import Filter, Sort, MetadataQuery from weaviate.collections.classes.data import DataObject from weaviate.classes.aggregate import GroupByAggregate -from weaviate.classes.init import AdditionalConfig, Timeout +from weaviate.classes.init import Auth, AdditionalConfig, Timeout +from weaviate.classes.config import Property, DataType import os import asyncio @@ -16,239 +40,150 @@ from sklearn.decomposition import PCA - from goldenverba.components.document import Document -from goldenverba.components.interfaces import ( - Reader, - Chunker, - Embedding, - Retriever, - Generator, -) -from goldenverba.server.helpers import LoggerManager -from goldenverba.server.types import FileConfig, FileStatus - -# Import Readers -from goldenverba.components.reader.BasicReader import BasicReader -from goldenverba.components.reader.GitReader import GitReader -from goldenverba.components.reader.UnstructuredAPI import UnstructuredReader -from goldenverba.components.reader.AssemblyAIAPI import AssemblyAIReader -from goldenverba.components.reader.HTMLReader import HTMLReader -from goldenverba.components.reader.FirecrawlReader import FirecrawlReader -from goldenverba.components.reader.UpstageDocumentParse import ( - UpstageDocumentParseReader, -) - -# Import Chunkers -from goldenverba.components.chunking.TokenChunker import TokenChunker -from goldenverba.components.chunking.SentenceChunker import SentenceChunker -from goldenverba.components.chunking.RecursiveChunker import RecursiveChunker -from goldenverba.components.chunking.HTMLChunker import HTMLChunker -from goldenverba.components.chunking.MarkdownChunker import MarkdownChunker -from goldenverba.components.chunking.CodeChunker import CodeChunker -from goldenverba.components.chunking.JSONChunker import JSONChunker -from goldenverba.components.chunking.SemanticChunker import SemanticChunker - -# Import Embedders -from goldenverba.components.embedding.OpenAIEmbedder import OpenAIEmbedder -from goldenverba.components.embedding.CohereEmbedder import CohereEmbedder -from goldenverba.components.embedding.OllamaEmbedder import OllamaEmbedder -from goldenverba.components.embedding.UpstageEmbedder import UpstageEmbedder -from goldenverba.components.embedding.WeaviateEmbedder import WeaviateEmbedder -from goldenverba.components.embedding.VoyageAIEmbedder import VoyageAIEmbedder -from goldenverba.components.embedding.SentenceTransformersEmbedder import ( - SentenceTransformersEmbedder, -) - -# Import Retrievers -from goldenverba.components.retriever.WindowRetriever import WindowRetriever - -# Import Generators -from goldenverba.components.generation.CohereGenerator import CohereGenerator -from goldenverba.components.generation.AnthrophicGenerator import AnthropicGenerator -from goldenverba.components.generation.OllamaGenerator import OllamaGenerator -from goldenverba.components.generation.OpenAIGenerator import OpenAIGenerator -from goldenverba.components.generation.GroqGenerator import GroqGenerator -from goldenverba.components.generation.NovitaGenerator import NovitaGenerator -from goldenverba.components.generation.UpstageGenerator import UpstageGenerator - -try: - import tiktoken -except Exception: - msg.warn("tiktoken not installed, your base installation might be corrupted.") - -### Add new components here ### - -production = os.getenv("VERBA_PRODUCTION") -if production != "Production": - readers = [ - BasicReader(), - HTMLReader(), - GitReader(), - UnstructuredReader(), - AssemblyAIReader(), - FirecrawlReader(), - UpstageDocumentParseReader(), - ] - chunkers = [ - TokenChunker(), - SentenceChunker(), - RecursiveChunker(), - SemanticChunker(), - HTMLChunker(), - MarkdownChunker(), - CodeChunker(), - JSONChunker(), - ] - embedders = [ - OllamaEmbedder(), - SentenceTransformersEmbedder(), - WeaviateEmbedder(), - UpstageEmbedder(), - VoyageAIEmbedder(), - CohereEmbedder(), - OpenAIEmbedder(), - ] - retrievers = [WindowRetriever()] - generators = [ - OllamaGenerator(), - OpenAIGenerator(), - AnthropicGenerator(), - CohereGenerator(), - GroqGenerator(), - NovitaGenerator(), - UpstageGenerator(), - ] -else: - readers = [ - BasicReader(), - HTMLReader(), - GitReader(), - UnstructuredReader(), - AssemblyAIReader(), - FirecrawlReader(), - UpstageDocumentParseReader(), + + +class WeaviateManager: + """Low-level Weaviate client wrapper. See module docstring for full details.""" + + # Shared timeout config applied to all connection types. + # stream=300 covers long-running streaming generation responses. + _TIMEOUT = AdditionalConfig( + timeout=Timeout(init=60, query=300, insert=300, stream=300) + ) + + # Property schemas for each named collection. + # Defined here so verify_collection can create collections with the correct + # schema on first use — without needing callers to pass properties manually. + _DOCUMENT_PROPERTIES = [ + Property(name="title", data_type=DataType.TEXT), + Property(name="content", data_type=DataType.TEXT), + Property(name="extension", data_type=DataType.TEXT), + Property(name="fileSize", data_type=DataType.NUMBER), + Property(name="labels", data_type=DataType.TEXT_ARRAY), + Property(name="source", data_type=DataType.TEXT), + Property(name="meta", data_type=DataType.TEXT), + Property(name="metadata", data_type=DataType.TEXT), ] - chunkers = [ - TokenChunker(), - SentenceChunker(), - RecursiveChunker(), - SemanticChunker(), - HTMLChunker(), - MarkdownChunker(), - CodeChunker(), - JSONChunker(), + _CONFIG_PROPERTIES = [ + Property(name="config", data_type=DataType.TEXT), ] - embedders = [ - WeaviateEmbedder(), - VoyageAIEmbedder(), - UpstageEmbedder(), - CohereEmbedder(), - OpenAIEmbedder(), + _SUGGESTION_PROPERTIES = [ + Property(name="query", data_type=DataType.TEXT), + Property(name="timestamp", data_type=DataType.TEXT), ] - retrievers = [WindowRetriever()] - generators = [ - OpenAIGenerator(), - AnthropicGenerator(), - CohereGenerator(), - UpstageGenerator(), + _CHUNK_PROPERTIES = [ + Property(name="content", data_type=DataType.TEXT), + Property(name="chunk_id", data_type=DataType.INT), + Property(name="doc_uuid", data_type=DataType.TEXT), + Property(name="title", data_type=DataType.TEXT), + Property(name="pca", data_type=DataType.NUMBER_ARRAY), + Property(name="start_i", data_type=DataType.INT), + Property(name="end_i", data_type=DataType.INT), + Property(name="content_without_overlap", data_type=DataType.TEXT), + Property(name="labels", data_type=DataType.TEXT_ARRAY), ] - -### ----------------------- ### - - -class WeaviateManager: def __init__(self): self.document_collection_name = "VERBA_DOCUMENTS" self.config_collection_name = "VERBA_CONFIGURATION" self.suggestion_collection_name = "VERBA_SUGGESTIONS" - self.embedding_table = {} - - ### Connection Handling - - async def connect_to_cluster(self, w_url, w_key): - if w_url is not None and w_key is not None: - msg.info(f"Connecting to Weaviate Cluster {w_url} with Auth") - return weaviate.use_async_with_weaviate_cloud( - cluster_url=w_url, - auth_credentials=AuthApiKey(w_key), - additional_config=AdditionalConfig( - timeout=Timeout(init=60, query=300, insert=300) - ), - ) - else: + + # Maps embedder model name → Weaviate collection name. + # Populated lazily on first use of each embedder. + self.embedding_table: dict[str, str] = {} + + # Separate table for cache collections so verify_cache_collection and + # verify_embedding_collection never clobber each other's entries. + self.cache_table: dict[str, str] = {} + + # Per-client set of collection names already confirmed to exist. + # Avoids a network round-trip on every operation; keyed by id(client). + self._verified_collections: dict[int, set[str]] = {} + + # ------------------------------------------------------------------------- + # Connection factories + # ------------------------------------------------------------------------- + + def connect_to_cluster(self, w_url: str, w_key: str) -> WeaviateAsyncClient: + """Connect to Weaviate Cloud. Both URL and API key are required.""" + if not w_url or not w_key: raise Exception("No URL or API Key provided") + msg.info(f"Connecting to Weaviate Cluster {w_url} with Auth") + return weaviate.use_async_with_weaviate_cloud( + cluster_url=w_url, + auth_credentials=Auth.api_key(w_key), + additional_config=self._TIMEOUT, + ) + + def connect_to_docker(self, host: str = "localhost") -> WeaviateAsyncClient: + """ + Connect to a local Weaviate Docker instance. - async def connect_to_docker(self, w_url): - msg.info(f"Connecting to Weaviate Docker") + Host defaults to localhost for local dev. In Docker Compose the service + name ("weaviate") is resolved by Docker DNS, so pass that explicitly. + Falls back to the WEAVIATE_HOST env var if set and no host is given. + """ + resolved_host = host or os.environ.get("WEAVIATE_HOST", "localhost") + msg.info(f"Connecting to Weaviate Docker at {resolved_host}") return weaviate.use_async_with_local( - host=w_url, - additional_config=AdditionalConfig( - timeout=Timeout(init=60, query=300, insert=300) - ), + host=resolved_host, + additional_config=self._TIMEOUT, ) - async def connect_to_custom(self, host, w_key, port): - # Extract the port from the host - msg.info(f"Connecting to Weaviate Custom") - - if host is None or host == "": + def connect_to_custom( + self, host: str, w_key: str, port: str + ) -> WeaviateAsyncClient: + """Connect to any reachable Weaviate instance with optional API key auth.""" + if not host: raise Exception("No Host URL provided") - - if w_key is None or w_key == "": - return weaviate.use_async_with_local( - host=host, - port=int(port), - skip_init_checks=True, - additional_config=AdditionalConfig( - timeout=Timeout(init=60, query=300, insert=300) - ), - ) - else: - return weaviate.use_async_with_local( - host=host, - port=int(port), - skip_init_checks=True, - auth_credentials=AuthApiKey(w_key), - additional_config=AdditionalConfig( - timeout=Timeout(init=60, query=300, insert=300) - ), - ) - - async def connect_to_embedded(self): - msg.info(f"Connecting to Weaviate Embedded") - return weaviate.use_async_with_embedded( - additional_config=AdditionalConfig( - timeout=Timeout(init=60, query=300, insert=300) - ) + msg.info(f"Connecting to Weaviate Custom at {host}:{port}") + kwargs = dict( + host=host, + port=int(port), + skip_init_checks=True, + additional_config=self._TIMEOUT, ) + if w_key: + kwargs["auth_credentials"] = Auth.api_key(w_key) + return weaviate.use_async_with_local(**kwargs) async def connect( self, deployment: str, weaviateURL: str, weaviateAPIKey: str, port: str = "8080" ) -> WeaviateAsyncClient: - try: + """ + Create and connect a WeaviateAsyncClient for the given deployment type. + + Supported deployments: + "Weaviate" — Weaviate Cloud (WEAVIATE_URL_VERBA / WEAVIATE_API_KEY_VERBA) + "Docker" — local Docker instance (WEAVIATE_HOST or localhost:8080) + "Custom" — user-supplied host + optional API key + Raises on any connection failure so ClientManager can surface the error. + """ + try: if deployment == "Weaviate": - if weaviateURL == "" and os.environ.get("WEAVIATE_URL_VERBA"): - weaviateURL = os.environ.get("WEAVIATE_URL_VERBA") - if weaviateAPIKey == "" and os.environ.get("WEAVIATE_API_KEY_VERBA"): - weaviateAPIKey = os.environ.get("WEAVIATE_API_KEY_VERBA") - client = await self.connect_to_cluster(weaviateURL, weaviateAPIKey) + # Fall back to env vars if the frontend sent empty strings. + weaviateURL = weaviateURL or os.environ.get("WEAVIATE_URL_VERBA", "") + weaviateAPIKey = weaviateAPIKey or os.environ.get( + "WEAVIATE_API_KEY_VERBA", "" + ) + client = self.connect_to_cluster(weaviateURL, weaviateAPIKey) elif deployment == "Docker": - client = await self.connect_to_docker("weaviate") - elif deployment == "Local": - client = await self.connect_to_embedded() + # In Docker Compose the Weaviate container is reachable as "weaviate". + # For local dev (outside Compose) localhost is correct. + client = self.connect_to_docker( + os.environ.get("WEAVIATE_HOST", "weaviate") + ) elif deployment == "Custom": - client = await self.connect_to_custom(weaviateURL, weaviateAPIKey, port) + client = self.connect_to_custom(weaviateURL, weaviateAPIKey, port) else: - raise Exception(f"Invalid deployment type: {deployment}") + raise Exception(f"Invalid deployment type: {deployment!r}") - if client is not None: - await client.connect() - if await client.is_ready(): - msg.good("Succesfully Connected to Weaviate") - return client + await client.connect() + if await client.is_ready(): + msg.good("Successfully Connected to Weaviate") + return client return None @@ -259,17 +194,23 @@ async def connect( ) async def disconnect(self, client: WeaviateAsyncClient): + """Close the client connection and free gRPC resources.""" try: await client.close() + # Discard the per-client collection cache so a future reconnect + # starts fresh (collections may have changed while disconnected). + self._verified_collections.pop(id(client), None) return True except Exception as e: msg.fail(f"Couldn't disconnect Weaviate: {str(e)}") return False - ### Metadata + # ------------------------------------------------------------------------- + # Cluster metadata + # ------------------------------------------------------------------------- async def get_metadata(self, client: WeaviateAsyncClient): - + """Return node info and per-collection object counts for the status page.""" # Node Information nodes = await client.cluster.nodes(output="verbose") node_payload = {"node_count": 0, "weaviate_version": "", "nodes": []} @@ -286,11 +227,10 @@ async def get_metadata(self, client: WeaviateAsyncClient): node_payload["weaviate_version"] = nodes[0].version # Collection Information - collections = await client.collections.list_all() collection_payload = {"collection_count": 0, "collections": []} for collection_name in collections: - collection_objects = await client.collections.get(collection_name).length() + collection_objects = await client.collections.use(collection_name).length() collection_payload["collections"].append( {"name": collection_name, "count": collection_objects} ) @@ -299,24 +239,55 @@ async def get_metadata(self, client: WeaviateAsyncClient): return node_payload, collection_payload - ### Collection Handling + # ------------------------------------------------------------------------- + # Collection management + # ------------------------------------------------------------------------- async def verify_collection( self, client: WeaviateAsyncClient, collection_name: str - ): + ) -> bool: + """ + Ensure a collection exists with the correct schema, creating it if needed. + + Results are cached per client so the exists() network call only happens + once per collection per connection lifetime — not on every operation. + + The schema (property list) is determined automatically from the collection + name so callers never need to pass schema details. + """ + client_id = id(client) + verified = self._verified_collections.setdefault(client_id, set()) + if collection_name in verified: + return True + if not await client.collections.exists(collection_name): - msg.info( - f"Collection: {collection_name} does not exist, creating new collection." + msg.info(f"Collection {collection_name!r} does not exist, creating it.") + properties = self._schema_for(collection_name) + await client.collections.create( + name=collection_name, + properties=properties, ) - returned_collection = await client.collections.create(name=collection_name) - if returned_collection: - return True - else: - return False - else: - return True - async def verify_embedding_collection(self, client: WeaviateAsyncClient, embedder): + verified.add(collection_name) + return True + + def _schema_for(self, collection_name: str) -> list[Property]: + """Return the property list for a given collection name.""" + if collection_name == self.document_collection_name: + return self._DOCUMENT_PROPERTIES + if collection_name == self.config_collection_name: + return self._CONFIG_PROPERTIES + if collection_name == self.suggestion_collection_name: + return self._SUGGESTION_PROPERTIES + if collection_name.startswith("VERBA_Embedding_"): + return self._CHUNK_PROPERTIES + # Cache collections and any unknown names: no predefined properties. + return [] + + async def verify_embedding_collection( + self, client: WeaviateAsyncClient, embedder: str + ) -> bool: + """Ensure the chunk+vector collection for `embedder` exists, creating it if needed.""" if embedder not in self.embedding_table: self.embedding_table[embedder] = "VERBA_Embedding_" + re.sub( r"[^a-zA-Z0-9]", "_", embedder @@ -325,18 +296,25 @@ async def verify_embedding_collection(self, client: WeaviateAsyncClient, embedde else: return True - async def verify_cache_collection(self, client: WeaviateAsyncClient, embedder): - if embedder not in self.embedding_table: - self.embedding_table[embedder] = "VERBA_Cache_" + re.sub( + async def verify_cache_collection( + self, client: WeaviateAsyncClient, embedder: str + ) -> bool: + """Ensure the semantic-cache collection for `embedder` exists, creating it if needed.""" + # Use a separate cache_table so this never collides with embedding_table entries + if embedder not in self.cache_table: + self.cache_table[embedder] = "VERBA_Cache_" + re.sub( r"[^a-zA-Z0-9]", "_", embedder ) - return await self.verify_collection(client, self.embedding_table[embedder]) + return await self.verify_collection(client, self.cache_table[embedder]) else: return True async def verify_embedding_collections( self, client: WeaviateAsyncClient, environment_variables, libraries ): + # Import here to avoid circular imports (managers.py defines `embedders` at module level) + from goldenverba.components.embedding.embedding_manager import embedders + for embedder in embedders: if embedder.check_available(environment_variables, libraries): if "Model" in embedder.config: @@ -359,11 +337,13 @@ async def verify_collections( ) return True - ### Configuration Handling + # ------------------------------------------------------------------------- + # Configuration storage (RAG config, theme, user prefs) + # ------------------------------------------------------------------------- async def get_config(self, client: WeaviateAsyncClient, uuid: str) -> dict: if await self.verify_collection(client, self.config_collection_name): - config_collection = client.collections.get(self.config_collection_name) + config_collection = client.collections.use(self.config_collection_name) if await config_collection.data.exists(uuid): config = await config_collection.query.fetch_object_by_id(uuid) return json.loads(config.properties["config"]) @@ -372,7 +352,7 @@ async def get_config(self, client: WeaviateAsyncClient, uuid: str) -> dict: async def set_config(self, client: WeaviateAsyncClient, uuid: str, config: dict): if await self.verify_collection(client, self.config_collection_name): - config_collection = client.collections.get(self.config_collection_name) + config_collection = client.collections.use(self.config_collection_name) if await config_collection.data.exists(uuid): if await config_collection.data.delete_by_id(uuid): await config_collection.data.insert( @@ -385,20 +365,29 @@ async def set_config(self, client: WeaviateAsyncClient, uuid: str, config: dict) async def reset_config(self, client: WeaviateAsyncClient, uuid: str): if await self.verify_collection(client, self.config_collection_name): - config_collection = client.collections.get(self.config_collection_name) + config_collection = client.collections.use(self.config_collection_name) if await config_collection.data.exists(uuid): await config_collection.data.delete_by_id(uuid) - ### Import Handling + # ------------------------------------------------------------------------- + # Document import + # ------------------------------------------------------------------------- async def import_document( self, client: WeaviateAsyncClient, document: Document, embedder: str ): + """ + Write a vectorized Document to Weaviate. + + Inserts the document record, then batch-inserts all chunks with their + vectors. Verifies chunk count after insertion and rolls back both the + document and chunks if there's a mismatch. + """ if await self.verify_collection( client, self.document_collection_name ) and await self.verify_embedding_collection(client, embedder): - document_collection = client.collections.get(self.document_collection_name) - embedder_collection = client.collections.get(self.embedding_table[embedder]) + document_collection = client.collections.use(self.document_collection_name) + embedder_collection = client.collections.use(self.embedding_table[embedder]) ### Import Document document_obj = Document.to_json(document) @@ -445,40 +434,53 @@ async def import_document( await self.delete_document(client, doc_uuid) raise Exception(f"Chunk import failed with : {str(e)}") - ### Document CRUD - - async def exist_document_name(self, client: WeaviateAsyncClient, name: str) -> str: - if await self.verify_collection(client, self.document_collection_name): - document_collection = client.collections.get(self.document_collection_name) - aggregation = await document_collection.aggregate.over_all(total_count=True) + # ------------------------------------------------------------------------- + # Document CRUD + # ------------------------------------------------------------------------- - if aggregation.total_count == 0: - return None - else: - documents = await document_collection.query.fetch_objects( - filters=Filter.by_property("title").equal(name) - ) - if len(documents.objects) > 0: - return documents.objects[0].uuid + async def exist_document_name( + self, client: WeaviateAsyncClient, name: str + ) -> str | None: + """ + Return the UUID of an existing document with the given title, or None. + A single filtered query suffices — the prior aggregate.over_all() that + checked total_count was an extra round-trip that gained nothing. + """ + if await self.verify_collection(client, self.document_collection_name): + document_collection = client.collections.use(self.document_collection_name) + documents = await document_collection.query.fetch_objects( + filters=Filter.by_property("title").equal(name), + limit=1, + ) + if documents.objects: + return documents.objects[0].uuid return None async def delete_document(self, client: WeaviateAsyncClient, uuid: str): if await self.verify_collection(client, self.document_collection_name): - document_collection = client.collections.get(self.document_collection_name) + document_collection = client.collections.use(self.document_collection_name) if not await document_collection.data.exists(uuid): return document_obj = await document_collection.query.fetch_object_by_id(uuid) - embedding_config = json.loads(document_obj.properties.get("meta"))[ - "Embedder" - ] - embedder = embedding_config["config"]["Model"]["value"] + meta_raw = document_obj.properties.get("meta") + if not meta_raw: + await document_collection.data.delete_by_id(uuid) + return + try: + embedding_config = json.loads(meta_raw)["Embedder"] + embedder = embedding_config["config"]["Model"]["value"] + except (json.JSONDecodeError, KeyError): + # meta is malformed or missing Embedder key — delete the document + # record but skip trying to clean up chunks we can't identify + await document_collection.data.delete_by_id(uuid) + return if await self.verify_embedding_collection(client, embedder): if await document_collection.data.delete_by_id(uuid): - embedder_collection = client.collections.get( + embedder_collection = client.collections.use( self.embedding_table[embedder] ) await embedder_collection.data.delete_many( @@ -486,14 +488,25 @@ async def delete_document(self, client: WeaviateAsyncClient, uuid: str): ) async def delete_all_documents(self, client: WeaviateAsyncClient): + """ + Delete all documents and their associated chunks. + + Collects all UUIDs first, then fires all deletes concurrently via + asyncio.gather() instead of processing them one by one. + """ if await self.verify_collection(client, self.document_collection_name): - document_collection = client.collections.get(self.document_collection_name) - async for item in document_collection.iterator(): - await self.delete_document(client, item.uuid) + document_collection = client.collections.use(self.document_collection_name) + all_uuids = [ + str(item.uuid) async for item in document_collection.iterator() + ] + await asyncio.gather( + *[self.delete_document(client, uuid) for uuid in all_uuids], + return_exceptions=True, + ) async def delete_all_configs(self, client: WeaviateAsyncClient): if await self.verify_collection(client, self.config_collection_name): - config_collection = client.collections.get(self.config_collection_name) + config_collection = client.collections.use(self.config_collection_name) async for item in config_collection.iterator(): await config_collection.data.delete_by_id(item.uuid) @@ -514,7 +527,7 @@ async def get_documents( ) -> list[dict]: if await self.verify_collection(client, self.document_collection_name): offset = pageSize * (page - 1) - document_collection = client.collections.get(self.document_collection_name) + document_collection = client.collections.use(self.document_collection_name) if len(labels) > 0: filter = Filter.by_property("labels").contains_all(labels) @@ -561,7 +574,7 @@ async def get_document( self, client: WeaviateAsyncClient, uuid: str, properties: list[str] = None ) -> list[dict]: if await self.verify_collection(client, self.document_collection_name): - document_collection = client.collections.get(self.document_collection_name) + document_collection = client.collections.use(self.document_collection_name) if await document_collection.data.exists(uuid): response = await document_collection.query.fetch_object_by_id( @@ -572,11 +585,13 @@ async def get_document( msg.warn(f"Document not found ({uuid})") return None - ### Labels + # ------------------------------------------------------------------------- + # Labels + # ------------------------------------------------------------------------- async def get_labels(self, client: WeaviateAsyncClient) -> list[str]: if await self.verify_collection(client, self.document_collection_name): - document_collection = client.collections.get(self.document_collection_name) + document_collection = client.collections.use(self.document_collection_name) aggregation = await document_collection.aggregate.over_all( group_by=GroupByAggregate(prop="labels"), total_count=True ) @@ -585,13 +600,15 @@ async def get_labels(self, client: WeaviateAsyncClient) -> list[str]: for aggregation_group in aggregation.groups ] - ### Chunks Retrieval + # ------------------------------------------------------------------------- + # Chunk retrieval + # ------------------------------------------------------------------------- async def get_chunk( self, client: WeaviateAsyncClient, uuid: str, embedder: str ) -> list[dict]: if await self.verify_embedding_collection(client, embedder): - embedder_collection = client.collections.get(self.embedding_table[embedder]) + embedder_collection = client.collections.use(self.embedding_table[embedder]) if await embedder_collection.data.exists(uuid): response = await embedder_collection.query.fetch_object_by_id(uuid) response.properties["doc_uuid"] = str(response.properties["doc_uuid"]) @@ -615,7 +632,7 @@ async def get_chunks( embedder = embedding_config["config"]["Model"]["value"] if await self.verify_embedding_collection(client, embedder): - embedder_collection = client.collections.get( + embedder_collection = client.collections.use( self.embedding_table[embedder] ) @@ -643,17 +660,14 @@ async def get_vectors( embedder = embedding_config["config"]["Model"]["value"] if await self.verify_embedding_collection(client, embedder): - embedder_collection = client.collections.get(self.embedding_table[embedder]) + embedder_collection = client.collections.use(self.embedding_table[embedder]) if not showAll: batch_size = 250 all_chunks = [] offset = 0 - total_time = 0 - call_count = 0 while True: - call_start_time = asyncio.get_event_loop().time() weaviate_chunks = await embedder_collection.query.fetch_objects( filters=Filter.by_property("doc_uuid").equal(uuid), limit=batch_size, @@ -661,10 +675,6 @@ async def get_vectors( return_properties=["chunk_id", "pca"], include_vector=True, ) - call_end_time = asyncio.get_event_loop().time() - call_duration = call_end_time - call_start_time - total_time += call_duration - call_count += 1 all_chunks.extend(weaviate_chunks.objects) @@ -692,31 +702,48 @@ async def get_vectors( # Generate PCA for all embeddings else: - vector_map = {} + # First pass: stream all items into memory + all_items = [] + dimensions = 0 + async for item in embedder_collection.iterator(include_vector=True): + all_items.append(item) + dimensions = len(item.vector["default"]) + + if not all_items: + return {"embedder": embedder, "dimensions": 0, "groups": []} + + # Batch-fetch all unique documents concurrently instead of one + # sequential get_document() call per unique doc_uuid inside the loop + unique_doc_uuids = list( + {str(item.properties["doc_uuid"]) for item in all_items} + ) + doc_results = await asyncio.gather( + *[ + self.get_document(client, doc_uuid, properties=["title"]) + for doc_uuid in unique_doc_uuids + ], + return_exceptions=True, + ) + vector_map = { + doc_uuid: {"name": doc["title"], "chunks": []} + for doc_uuid, doc in zip(unique_doc_uuids, doc_results) + if doc and not isinstance(doc, Exception) + } + + # Second pass: collect vectors for successfully-fetched documents vector_list, vector_ids, vector_chunk_uuids, vector_chunk_ids = ( [], [], [], [], ) - dimensions = 0 - - async for item in embedder_collection.iterator(include_vector=True): - doc_uuid = item.properties["doc_uuid"] - chunk_uuid = item.uuid + for item in all_items: + doc_uuid = str(item.properties["doc_uuid"]) if doc_uuid not in vector_map: - _document = await self.get_document(client, doc_uuid) - if _document: - vector_map[doc_uuid] = { - "name": _document["title"], - "chunks": [], - } - else: - continue + continue vector_list.append(item.vector["default"]) - dimensions = len(item.vector["default"]) vector_ids.append(doc_uuid) - vector_chunk_uuids.append(chunk_uuid) + vector_chunk_uuids.append(item.uuid) vector_chunk_ids.append(item.properties["chunk_id"]) if len(vector_ids) > 3: @@ -758,6 +785,25 @@ async def get_vectors( return None + async def get_chunk_by_ids( + self, client: WeaviateAsyncClient, embedder: str, doc_uuid: str, ids: list[int] + ): + """Fetch specific chunks by their sequential chunk_id values within a document.""" + if await self.verify_embedding_collection(client, embedder): + embedder_collection = client.collections.use(self.embedding_table[embedder]) + try: + weaviate_chunks = await embedder_collection.query.fetch_objects( + filters=( + Filter.by_property("doc_uuid").equal(str(doc_uuid)) + & Filter.by_property("chunk_id").contains_any(list(ids)) + ), + sort=Sort.by_property("chunk_id", ascending=True), + ) + return weaviate_chunks.objects + except Exception as e: + msg.fail(f"Failed to fetch chunks: {str(e)}") + raise e + async def hybrid_chunks( self, client: WeaviateAsyncClient, @@ -769,8 +815,15 @@ async def hybrid_chunks( labels: list[str], document_uuids: list[str], ): + """ + Run a hybrid (BM25 + vector) search over the embedder's chunk collection. + + limit_mode="Autocut" uses Weaviate's auto_limit (stops at a natural + relevance cutoff); otherwise a hard limit is applied. + Filters are ANDed together — labels AND document_uuids if both are given. + """ if await self.verify_embedding_collection(client, embedder): - embedder_collection = client.collections.get(self.embedding_table[embedder]) + embedder_collection = client.collections.use(self.embedding_table[embedder]) filters = [] @@ -810,42 +863,23 @@ async def hybrid_chunks( return chunks.objects - async def get_chunk_by_ids( - self, client: WeaviateAsyncClient, embedder: str, doc_uuid: str, ids: list[int] - ): - if await self.verify_embedding_collection(client, embedder): - embedder_collection = client.collections.get(self.embedding_table[embedder]) - try: - weaviate_chunks = await embedder_collection.query.fetch_objects( - filters=( - Filter.by_property("doc_uuid").equal(str(doc_uuid)) - & Filter.by_property("chunk_id").contains_any(list(ids)) - ), - sort=Sort.by_property("chunk_id", ascending=True), - ) - return weaviate_chunks.objects - except Exception as e: - msg.fail(f"Failed to fetch chunks: {str(e)}") - raise e - - ### Suggestion Logic + # ------------------------------------------------------------------------- + # Suggestions + # ------------------------------------------------------------------------- async def add_suggestion(self, client: WeaviateAsyncClient, query: str): + """Store a query as an autocomplete suggestion (deduplicates by exact match).""" if await self.verify_collection(client, self.suggestion_collection_name): - suggestion_collection = client.collections.get( + suggestion_collection = client.collections.use( self.suggestion_collection_name ) - aggregation = await suggestion_collection.aggregate.over_all( - total_count=True + # One query suffices — if the collection is empty the filter returns 0 + # results regardless, so a prior aggregate.over_all() is redundant. + existing = await suggestion_collection.query.fetch_objects( + filters=Filter.by_property("query").equal(query) ) - if aggregation.total_count > 0: - does_suggestion_exists = ( - await suggestion_collection.query.fetch_objects( - filters=Filter.by_property("query").equal(query) - ) - ) - if len(does_suggestion_exists.objects) > 0: - return + if len(existing.objects) > 0: + return await suggestion_collection.data.insert( {"query": query, "timestamp": datetime.now().isoformat()} ) @@ -854,7 +888,7 @@ async def retrieve_suggestions( self, client: WeaviateAsyncClient, query: str, limit: int ): if await self.verify_collection(client, self.suggestion_collection_name): - suggestion_collection = client.collections.get( + suggestion_collection = client.collections.use( self.suggestion_collection_name ) suggestions = await suggestion_collection.query.bm25( @@ -874,7 +908,7 @@ async def retrieve_all_suggestions( self, client: WeaviateAsyncClient, page: int, pageSize: int ): if await self.verify_collection(client, self.suggestion_collection_name): - suggestion_collection = client.collections.get( + suggestion_collection = client.collections.use( self.suggestion_collection_name ) offset = pageSize * (page - 1) @@ -898,7 +932,7 @@ async def retrieve_all_suggestions( async def delete_suggestions(self, client: WeaviateAsyncClient, uuid: str): if await self.verify_collection(client, self.suggestion_collection_name): - suggestion_collection = client.collections.get( + suggestion_collection = client.collections.use( self.suggestion_collection_name ) await suggestion_collection.data.delete_by_id(uuid) @@ -907,17 +941,15 @@ async def delete_all_suggestions(self, client: WeaviateAsyncClient): if await self.verify_collection(client, self.suggestion_collection_name): await client.collections.delete(self.suggestion_collection_name) - ### Cache Logic - - # TODO: Implement Cache Logic - - ### Metadata Retrieval + # ------------------------------------------------------------------------- + # Metadata / counts + # ------------------------------------------------------------------------- async def get_datacount( self, client: WeaviateAsyncClient, embedder: str, document_uuids: list[str] = [] ) -> int: if await self.verify_embedding_collection(client, embedder): - embedder_collection = client.collections.get(self.embedding_table[embedder]) + embedder_collection = client.collections.use(self.embedding_table[embedder]) if document_uuids: filters = Filter.by_property("doc_uuid").contains_any(document_uuids) @@ -938,7 +970,7 @@ async def get_chunk_count( self, client: WeaviateAsyncClient, embedder: str, doc_uuid: str ) -> int: if await self.verify_embedding_collection(client, embedder): - embedder_collection = client.collections.get(self.embedding_table[embedder]) + embedder_collection = client.collections.use(self.embedding_table[embedder]) response = await embedder_collection.aggregate.over_all( filters=Filter.by_property("doc_uuid").equal(doc_uuid), group_by=GroupByAggregate(prop="doc_uuid"), @@ -948,339 +980,3 @@ async def get_chunk_count( return response.groups[0].total_count else: return 0 - - -class ReaderManager: - def __init__(self): - self.readers: dict[str, Reader] = {reader.name: reader for reader in readers} - - async def load( - self, reader: str, fileConfig: FileConfig, logger: LoggerManager - ) -> list[Document]: - try: - loop = asyncio.get_running_loop() - start_time = loop.time() - if reader in self.readers: - config = fileConfig.rag_config["Reader"].components[reader].config - documents: list[Document] = await self.readers[reader].load( - config, fileConfig - ) - for document in documents: - document.meta["Reader"] = ( - fileConfig.rag_config["Reader"].components[reader].model_dump() - ) - elapsed_time = round(loop.time() - start_time, 2) - if len(documents) == 1: - await logger.send_report( - fileConfig.fileID, - FileStatus.LOADING, - f"Loaded {fileConfig.filename}", - took=elapsed_time, - ) - else: - await logger.send_report( - fileConfig.fileID, - FileStatus.LOADING, - f"Loaded {fileConfig.filename} with {len(documents)} documents", - took=elapsed_time, - ) - await logger.send_report( - fileConfig.fileID, FileStatus.CHUNKING, "", took=0 - ) - return documents - else: - raise Exception(f"{reader} Reader not found") - - except Exception as e: - raise Exception(f"Reader {reader} failed with: {str(e)}") - - -class ChunkerManager: - def __init__(self): - self.chunkers: dict[str, Chunker] = { - chunker.name: chunker for chunker in chunkers - } - - async def chunk( - self, - chunker: str, - fileConfig: FileConfig, - documents: list[Document], - embedder: Embedding, - logger: LoggerManager, - ) -> list[Document]: - try: - loop = asyncio.get_running_loop() - start_time = loop.time() - if chunker in self.chunkers: - config = fileConfig.rag_config["Chunker"].components[chunker].config - embedder_config = ( - fileConfig.rag_config["Embedder"].components[embedder.name].config - ) - chunked_documents = await self.chunkers[chunker].chunk( - config=config, - documents=documents, - embedder=embedder, - embedder_config=embedder_config, - ) - for chunked_document in chunked_documents: - chunked_document.meta["Chunker"] = ( - fileConfig.rag_config["Chunker"] - .components[chunker] - .model_dump() - ) - elapsed_time = round(loop.time() - start_time, 2) - if len(documents) == 1: - await logger.send_report( - fileConfig.fileID, - FileStatus.CHUNKING, - f"Split {fileConfig.filename} into {len(chunked_documents[0].chunks)} chunks", - took=elapsed_time, - ) - else: - await logger.send_report( - fileConfig.fileID, - FileStatus.CHUNKING, - f"Chunked all {len(chunked_documents)} documents with a total of {sum([len(document.chunks) for document in chunked_documents])} chunks", - took=elapsed_time, - ) - - await logger.send_report( - fileConfig.fileID, FileStatus.EMBEDDING, "", took=0 - ) - return chunked_documents - else: - raise Exception(f"{chunker} Chunker not found") - except Exception as e: - raise e - - -class EmbeddingManager: - def __init__(self): - self.embedders: dict[str, Embedding] = { - embedder.name: embedder for embedder in embedders - } - - async def vectorize( - self, - embedder: str, - fileConfig: FileConfig, - documents: list[Document], - logger: LoggerManager, - ) -> list[Document]: - """Vectorizes chunks in batches - @parameter: documents : Document - Verba document - @returns Document - Document with vectorized chunks - """ - try: - loop = asyncio.get_running_loop() - start_time = loop.time() - if embedder in self.embedders: - config = fileConfig.rag_config["Embedder"].components[embedder].config - - for document in documents: - content = [ - document.metadata + "\n" + chunk.content - for chunk in document.chunks - ] - embeddings = await self.batch_vectorize(embedder, config, content) - - if len(embeddings) >= 3: - pca = PCA(n_components=3) - generated_pca_embeddings = pca.fit_transform(embeddings) - pca_embeddings = [ - pca_.tolist() for pca_ in generated_pca_embeddings - ] - else: - pca_embeddings = [embedding[0:3] for embedding in embeddings] - - for vector, chunk, pca_ in zip( - embeddings, document.chunks, pca_embeddings - ): - chunk.vector = vector - chunk.pca = pca_ - - document.meta["Embedder"] = ( - fileConfig.rag_config["Embedder"] - .components[embedder] - .model_dump() - ) - - elapsed_time = round(loop.time() - start_time, 2) - await logger.send_report( - fileConfig.fileID, - FileStatus.EMBEDDING, - f"Vectorized all chunks", - took=elapsed_time, - ) - await logger.send_report( - fileConfig.fileID, FileStatus.INGESTING, "", took=0 - ) - return documents - else: - raise Exception(f"{embedder} Embedder not found") - except Exception as e: - raise e - - async def batch_vectorize( - self, embedder: str, config: dict, content: list[str] - ) -> list[list[float]]: - """Vectorize content in batches""" - try: - batches = [ - content[i : i + self.embedders[embedder].max_batch_size] - for i in range(0, len(content), self.embedders[embedder].max_batch_size) - ] - msg.info(f"Vectorizing {len(content)} chunks in {len(batches)} batches") - tasks = [ - self.embedders[embedder].vectorize(config, batch) for batch in batches - ] - results = await asyncio.gather(*tasks, return_exceptions=True) - - # Check if all tasks were successful - errors = [r for r in results if isinstance(r, Exception)] - if errors: - error_messages = [str(e) for e in errors] - raise Exception( - f"Vectorization failed for some batches: {', '.join(error_messages)}" - ) - - # Flatten the results - flattened_results = [item for sublist in results for item in sublist] - - # Verify the number of vectors matches the input content - if len(flattened_results) != len(content): - raise Exception( - f"Mismatch in vectorization results: expected {len(content)} vectors, got {len(flattened_results)}" - ) - - return flattened_results - except Exception as e: - raise Exception(f"Batch vectorization failed: {str(e)}") - - async def vectorize_query( - self, embedder: str, content: str, rag_config: dict - ) -> list[float]: - try: - if embedder in self.embedders: - config = rag_config["Embedder"].components[embedder].config - embeddings = await self.embedders[embedder].vectorize(config, [content]) - return embeddings[0] - else: - raise Exception(f"{embedder} Embedder not found") - except Exception as e: - raise e - - -class RetrieverManager: - def __init__(self): - self.retrievers: dict[str, Retriever] = { - retriever.name: retriever for retriever in retrievers - } - - async def retrieve( - self, - client, - retriever: str, - query: str, - vector: list[float], - rag_config: dict, - weaviate_manager: WeaviateManager, - labels: list[str], - document_uuids: list[str], - ): - try: - if retriever not in self.retrievers: - raise Exception(f"Retriever {retriever} not found") - - embedder_model = ( - rag_config["Embedder"] - .components[rag_config["Embedder"].selected] - .config["Model"] - .value - ) - config = rag_config["Retriever"].components[retriever].config - documents, context = await self.retrievers[retriever].retrieve( - client, - query, - vector, - config, - weaviate_manager, - embedder_model, - labels, - document_uuids, - ) - return (documents, context) - - except Exception as e: - raise e - - -class GeneratorManager: - def __init__(self): - self.generators: dict[str, Generator] = { - generator.name: generator for generator in generators - } - - async def generate_stream(self, rag_config, query, context, conversation): - """Generate a stream of response dicts based on a list of queries and list of contexts, and includes conversational context - @parameter: queries : list[str] - List of queries - @parameter: context : list[str] - List of contexts - @parameter: conversation : dict - Conversational context - @returns Iterator[dict] - Token response generated by the Generator in this format {system:TOKEN, finish_reason:stop or empty}. - """ - - generator = rag_config["Generator"].selected - generator_config = ( - rag_config["Generator"].components[rag_config["Generator"].selected].config - ) - - if generator not in self.generators: - raise Exception(f"Generator {generator} not found") - - async for result in self.generators[generator].generate_stream( - generator_config, query, context, conversation - ): - yield result - - def truncate_conversation_dicts( - self, conversation_dicts: list[dict[str, any]], max_tokens: int - ) -> list[dict[str, any]]: - """ - Truncate a list of conversation dictionaries to fit within a specified maximum token limit. - - @parameter conversation_dicts: List[Dict[str, any]] - A list of conversation dictionaries that may contain various keys, where 'content' key is present and contains text data. - @parameter max_tokens: int - The maximum number of tokens that the combined content of the truncated conversation dictionaries should not exceed. - - @returns List[Dict[str, any]]: A list of conversation dictionaries that have been truncated so that their combined content respects the max_tokens limit. The list is returned in the original order of conversation with the most recent conversation being truncated last if necessary. - - """ - encoding = tiktoken.encoding_for_model("gpt-3.5-turbo") - accumulated_tokens = 0 - truncated_conversation_dicts = [] - - # Start with the newest conversations - for item_dict in reversed(conversation_dicts): - item_tokens = encoding.encode(item_dict["content"], disallowed_special=()) - - # If adding the entire new item exceeds the max tokens - if accumulated_tokens + len(item_tokens) > max_tokens: - # Calculate how many tokens we can add from this item - remaining_space = max_tokens - accumulated_tokens - truncated_content = encoding.decode(item_tokens[:remaining_space]) - - # Create a new truncated item dictionary - truncated_item_dict = { - "type": item_dict["type"], - "content": truncated_content, - "typewriter": item_dict["typewriter"], - } - - truncated_conversation_dicts.append(truncated_item_dict) - break - - truncated_conversation_dicts.append(item_dict) - accumulated_tokens += len(item_tokens) - - # The list has been built in reverse order so we reverse it again - return list(reversed(truncated_conversation_dicts)) diff --git a/goldenverba/server/api.py b/goldenverba/server/api.py index 6c40cb1a84..a87c5160cf 100644 --- a/goldenverba/server/api.py +++ b/goldenverba/server/api.py @@ -15,7 +15,8 @@ from starlette.websockets import WebSocketDisconnect from wasabi import msg # type: ignore[import] -from goldenverba import verba_manager +from goldenverba.components import verba_manager +from goldenverba.components.client_manager import ClientManager from goldenverba.server.types import ( ResetPayload, @@ -54,14 +55,23 @@ manager = verba_manager.VerbaManager() -client_manager = verba_manager.ClientManager() +client_manager = ClientManager() ### Lifespan +async def _periodic_cleanup(): + """Run ClientManager cleanup every 5 minutes in the background.""" + while True: + await asyncio.sleep(300) + await client_manager.clean_up() + + @asynccontextmanager async def lifespan(app: FastAPI): + cleanup_task = asyncio.create_task(_periodic_cleanup()) yield + cleanup_task.cancel() await client_manager.disconnect() @@ -211,17 +221,14 @@ async def websocket_generate_stream(websocket: WebSocket): msg.good(f"Received generate stream call for {payload.query}") - full_text = "" async for chunk in manager.generate_stream_answer( payload.rag_config, payload.query, payload.context, payload.conversation, ): - full_text += chunk["message"] - if chunk["finish_reason"] == "stop": - chunk["full_text"] = full_text await websocket.send_json(chunk) + msg.good("Successfully streamed answer") except WebSocketDisconnect: msg.warn("WebSocket connection closed by client.") @@ -230,9 +237,8 @@ async def websocket_generate_stream(websocket: WebSocket): except Exception as e: msg.fail(f"WebSocket Error: {str(e)}") await websocket.send_json( - {"message": e, "finish_reason": "stop", "full_text": str(e)} + {"message": str(e), "finish_reason": "stop", "full_text": str(e)} ) - msg.good("Succesfully streamed answer") @app.websocket("/ws/import_files") diff --git a/goldenverba/server/cli.py b/goldenverba/server/cli.py index 266c051e79..4a926ac5bb 100644 --- a/goldenverba/server/cli.py +++ b/goldenverba/server/cli.py @@ -3,7 +3,7 @@ import os from dotenv import load_dotenv -from goldenverba import verba_manager +from goldenverba.components import verba_manager from goldenverba.server.types import Credentials load_dotenv() diff --git a/goldenverba/server/helpers.py b/goldenverba/server/helpers.py index 1854e75d54..8820171282 100644 --- a/goldenverba/server/helpers.py +++ b/goldenverba/server/helpers.py @@ -1,3 +1,5 @@ +import time + from fastapi import WebSocket from goldenverba.server.types import ( FileStatus, @@ -8,6 +10,8 @@ ) from wasabi import msg +_BATCH_TTL_SECONDS = 300 # abandon incomplete uploads after 5 minutes + class LoggerManager: def __init__(self, socket: WebSocket = None): @@ -45,15 +49,27 @@ class BatchManager: def __init__(self): self.batches = {} + def _evict_stale(self): + now = time.monotonic() + stale = [ + fid + for fid, entry in self.batches.items() + if now - entry["created_at"] > _BATCH_TTL_SECONDS + ] + for fid in stale: + msg.warn(f"Evicting stale upload {fid} from BatchManager (TTL exceeded)") + del self.batches[fid] + def add_batch(self, payload: DataBatchPayload) -> FileConfig: try: - # msg.info(f"Receiving Batch for {payload.fileID} : {payload.order} of {payload.total}") + self._evict_stale() if payload.fileID not in self.batches: self.batches[payload.fileID] = { "fileID": payload.fileID, "total": payload.total, "chunks": {}, + "created_at": time.monotonic(), } self.batches[payload.fileID]["chunks"][payload.order] = payload.chunk diff --git a/goldenverba/server/types.py b/goldenverba/server/types.py index 71dc1efbf6..082696c00d 100644 --- a/goldenverba/server/types.py +++ b/goldenverba/server/types.py @@ -1,5 +1,5 @@ -from typing import Literal -from pydantic import BaseModel +from typing import Annotated, Literal +from pydantic import BaseModel, Field from enum import Enum @@ -166,7 +166,7 @@ class DocumentFilter(BaseModel): class GetSuggestionsPayload(BaseModel): - query: str + query: Annotated[str, Field(max_length=50_000)] limit: int credentials: Credentials @@ -183,7 +183,7 @@ class GetAllSuggestionsPayload(BaseModel): class QueryPayload(BaseModel): - query: str + query: Annotated[str, Field(max_length=50_000)] RAG: dict[str, RAGComponentClass] labels: list[str] documentFilter: list[DocumentFilter] @@ -227,9 +227,9 @@ class GetContentPayload(BaseModel): class GeneratePayload(BaseModel): - query: str - context: str - conversation: list[ConversationItem] + query: Annotated[str, Field(max_length=50_000)] + context: Annotated[str, Field(max_length=500_000)] + conversation: Annotated[list[ConversationItem], Field(max_length=100)] rag_config: dict[str, RAGComponentClass] diff --git a/goldenverba/tests/__init__.py b/goldenverba/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/goldenverba/tests/components/__init__.py b/goldenverba/tests/components/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/goldenverba/tests/components/chunking/__init__.py b/goldenverba/tests/components/chunking/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/goldenverba/tests/components/chunking/test_chunkers.py b/goldenverba/tests/components/chunking/test_chunkers.py new file mode 100644 index 0000000000..95d952b3fd --- /dev/null +++ b/goldenverba/tests/components/chunking/test_chunkers.py @@ -0,0 +1,474 @@ +""" +Tests for the three built-in chunkers: + - TokenChunker (goldenverba/components/chunking/TokenChunker.py) + - SentenceChunker (goldenverba/components/chunking/SentenceChunker.py) + - MarkdownChunker (goldenverba/components/chunking/MarkdownChunker.py) + +These are pure text-processing operations – no network calls required. +""" +import pytest + +from goldenverba.components.document import Document +from goldenverba.components.chunk import Chunk +from goldenverba.components.types import InputConfig +from goldenverba.components.chunking.TokenChunker import TokenChunker +from goldenverba.components.chunking.SentenceChunker import SentenceChunker +from goldenverba.components.chunking.MarkdownChunker import MarkdownChunker + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + +SHORT_SENTENCE = "The quick brown fox jumps over the lazy dog." + +MULTI_SENTENCE = ( + "The sun rose early. Birds began to sing. The air was crisp and cool. " + "Children ran outside to play. A gentle breeze rustled the leaves. " + "It was a perfect morning. Everyone felt alive and refreshed." +) + +LONG_TOKEN_TEXT = " ".join([f"word{i}" for i in range(600)]) + +MARKDOWN_TEXT = """\ +# Introduction + +This is the introduction paragraph with some text. + +## Background + +Here is some background information that explains the context. + +### Details + +Fine-grained details live under this subsection. + +## Summary + +A brief wrap-up of what was covered. +""" + + +def _make_token_config(tokens: int = 50, overlap: int = 10) -> dict[str, InputConfig]: + return { + "Tokens": InputConfig( + type="number", value=tokens, description="Tokens per chunk", values=[] + ), + "Overlap": InputConfig( + type="number", value=overlap, description="Overlap tokens", values=[] + ), + } + + +def _make_sentence_config( + sentences: int = 2, overlap: int = 0 +) -> dict[str, InputConfig]: + return { + "Sentences": InputConfig( + type="number", + value=sentences, + description="Sentences per chunk", + values=[], + ), + "Overlap": InputConfig( + type="number", value=overlap, description="Overlap sentences", values=[] + ), + } + + +# --------------------------------------------------------------------------- +# TokenChunker +# --------------------------------------------------------------------------- + + +class TestTokenChunker: + def _chunker(self) -> TokenChunker: + return TokenChunker() + + @pytest.mark.asyncio + async def test_produces_chunks(self): + """TokenChunker should split a long document into multiple chunks.""" + doc = Document(content=LONG_TOKEN_TEXT) + result = await self._chunker().chunk(_make_token_config(50, 0), [doc]) + + assert len(result) == 1 + assert len(result[0].chunks) > 1, "Expected multiple chunks for a long text" + + @pytest.mark.asyncio + async def test_returns_list_of_chunk_objects(self): + """Every element in document.chunks must be a Chunk instance.""" + doc = Document(content=LONG_TOKEN_TEXT) + result = await self._chunker().chunk(_make_token_config(50, 0), [doc]) + + for chunk in result[0].chunks: + assert isinstance(chunk, Chunk) + + @pytest.mark.asyncio + async def test_chunks_have_content(self): + """Each chunk must have non-empty content.""" + doc = Document(content=LONG_TOKEN_TEXT) + result = await self._chunker().chunk(_make_token_config(50, 0), [doc]) + + for chunk in result[0].chunks: + assert chunk.content.strip() != "", "Chunk content must not be empty" + + @pytest.mark.asyncio + async def test_chunk_ids_are_sequential(self): + """chunk_id values must form a contiguous sequence starting at 0.""" + doc = Document(content=LONG_TOKEN_TEXT) + result = await self._chunker().chunk(_make_token_config(50, 0), [doc]) + + ids = [c.chunk_id for c in result[0].chunks] + assert ids == list(range(len(ids))) + + @pytest.mark.asyncio + async def test_respects_max_token_size(self): + """Each chunk should contain at most (tokens + overlap) words.""" + tokens = 30 + overlap = 5 + doc = Document(content=LONG_TOKEN_TEXT) + result = await self._chunker().chunk(_make_token_config(tokens, overlap), [doc]) + + for chunk in result[0].chunks: + word_count = len(chunk.content.split()) + assert word_count <= tokens + overlap + 5, ( + f"Chunk has {word_count} words, expected <= {tokens + overlap}" + ) + + @pytest.mark.asyncio + async def test_overlap_text_is_shared(self): + """With overlap > 0, the tail of one chunk should appear at the head of next.""" + tokens = 10 + overlap = 3 + doc = Document(content=LONG_TOKEN_TEXT) + result = await self._chunker().chunk(_make_token_config(tokens, overlap), [doc]) + chunks = result[0].chunks + + if len(chunks) >= 2: + # Last 'overlap' words of chunk 0 should appear in chunk 1 + tail_words = chunks[0].content.split()[-overlap:] + head_words = chunks[1].content.split()[:overlap] + assert tail_words == head_words, ( + "Overlap words from chunk[0] should appear at start of chunk[1]" + ) + + @pytest.mark.asyncio + async def test_empty_string_yields_one_chunk(self): + """An empty document should still produce exactly one (empty) chunk.""" + doc = Document(content="") + result = await self._chunker().chunk(_make_token_config(50, 0), [doc]) + + # The chunker either produces 0 chunks or 1 empty chunk; it must not crash + assert isinstance(result[0].chunks, list) + + @pytest.mark.asyncio + async def test_very_short_text_yields_single_chunk(self): + """Text shorter than one chunk boundary should result in a single chunk.""" + doc = Document(content=SHORT_SENTENCE) + result = await self._chunker().chunk(_make_token_config(250, 50), [doc]) + + assert len(result[0].chunks) == 1 + assert result[0].chunks[0].content == SHORT_SENTENCE + + @pytest.mark.asyncio + async def test_overlap_clamped_when_greater_than_units(self): + """When overlap >= units the chunker should clamp overlap and not crash.""" + doc = Document(content=LONG_TOKEN_TEXT) + # overlap (60) > tokens (50) -- the chunker should handle this gracefully + result = await self._chunker().chunk(_make_token_config(50, 60), [doc]) + + assert len(result[0].chunks) >= 1 + + @pytest.mark.asyncio + async def test_skips_already_chunked_documents(self): + """If a document already has chunks, the chunker must leave them unchanged.""" + doc = Document(content=LONG_TOKEN_TEXT) + pre_existing_chunk = Chunk( + content="pre-existing", + chunk_id=0, + start_i=0, + end_i=12, + content_without_overlap="pre-existing", + ) + doc.chunks.append(pre_existing_chunk) + + result = await self._chunker().chunk(_make_token_config(50, 0), [doc]) + + assert len(result[0].chunks) == 1 + assert result[0].chunks[0].content == "pre-existing" + + @pytest.mark.asyncio + async def test_content_without_overlap_shorter_or_equal_to_content(self): + """content_without_overlap must be <= content in length.""" + doc = Document(content=LONG_TOKEN_TEXT) + result = await self._chunker().chunk(_make_token_config(20, 5), [doc]) + + for chunk in result[0].chunks: + assert len(chunk.content_without_overlap) <= len(chunk.content) + + +# --------------------------------------------------------------------------- +# SentenceChunker +# --------------------------------------------------------------------------- + + +class TestSentenceChunker: + def _chunker(self) -> SentenceChunker: + return SentenceChunker() + + @pytest.mark.asyncio + async def test_splits_on_sentence_boundaries(self): + """SentenceChunker should produce more than one chunk for multi-sentence text.""" + doc = Document(content=MULTI_SENTENCE) + result = await self._chunker().chunk(_make_sentence_config(2, 0), [doc]) + + assert len(result[0].chunks) > 1 + + @pytest.mark.asyncio + async def test_returns_chunk_instances(self): + """Every element of document.chunks must be a Chunk.""" + doc = Document(content=MULTI_SENTENCE) + result = await self._chunker().chunk(_make_sentence_config(2, 0), [doc]) + + for chunk in result[0].chunks: + assert isinstance(chunk, Chunk) + + @pytest.mark.asyncio + async def test_chunk_content_is_nonempty(self): + """Chunks must not be blank.""" + doc = Document(content=MULTI_SENTENCE) + result = await self._chunker().chunk(_make_sentence_config(2, 0), [doc]) + + for chunk in result[0].chunks: + assert chunk.content.strip() != "" + + @pytest.mark.asyncio + async def test_chunk_ids_sequential(self): + """chunk_id values must be contiguous starting from 0.""" + doc = Document(content=MULTI_SENTENCE) + result = await self._chunker().chunk(_make_sentence_config(2, 0), [doc]) + ids = [c.chunk_id for c in result[0].chunks] + assert ids == list(range(len(ids))) + + @pytest.mark.asyncio + async def test_single_sentence_yields_one_chunk(self): + """A document with a single sentence should end up in a single chunk.""" + doc = Document(content=SHORT_SENTENCE) + result = await self._chunker().chunk(_make_sentence_config(5, 1), [doc]) + + assert len(result[0].chunks) == 1 + assert result[0].chunks[0].content.strip() == SHORT_SENTENCE.strip() + + @pytest.mark.asyncio + async def test_overlap_produces_shared_sentences(self): + """With sentence-level overlap the last sentence of one chunk should + appear at the start of the next chunk's content.""" + doc = Document(content=MULTI_SENTENCE) + # 2 sentences per chunk, 1 sentence overlap + result = await self._chunker().chunk(_make_sentence_config(2, 1), [doc]) + chunks = result[0].chunks + + if len(chunks) >= 2: + # The last sentence of chunk 0 must appear in chunk 1 + last_sentence_chunk0 = chunks[0].content.split(".")[-2].strip() + assert last_sentence_chunk0 in chunks[1].content, ( + "Overlapping sentence from chunk[0] should appear in chunk[1]" + ) + + @pytest.mark.asyncio + async def test_overlap_clamped_when_too_large(self): + """overlap >= sentences must be clamped without crashing.""" + doc = Document(content=MULTI_SENTENCE) + # overlap (5) >= sentences (3) + result = await self._chunker().chunk(_make_sentence_config(3, 5), [doc]) + assert len(result[0].chunks) >= 1 + + @pytest.mark.asyncio + async def test_skips_already_chunked_documents(self): + """Pre-chunked documents must not be re-chunked.""" + doc = Document(content=MULTI_SENTENCE) + sentinel = Chunk( + content="sentinel", + chunk_id=0, + start_i=0, + end_i=8, + content_without_overlap="sentinel", + ) + doc.chunks.append(sentinel) + + result = await self._chunker().chunk(_make_sentence_config(2, 0), [doc]) + + assert len(result[0].chunks) == 1 + assert result[0].chunks[0].content == "sentinel" + + @pytest.mark.asyncio + async def test_all_text_covered(self): + """Concatenating content_without_overlap for all chunks should cover + all of the original sentences (no sentence should be lost).""" + doc = Document(content=MULTI_SENTENCE) + result = await self._chunker().chunk(_make_sentence_config(2, 0), [doc]) + + combined = " ".join(c.content_without_overlap for c in result[0].chunks) + # Every sentence from the source must appear somewhere in the combination + for sentence in MULTI_SENTENCE.split("."): + sentence = sentence.strip() + if sentence: + assert sentence in combined, ( + f"Sentence '{sentence}' is missing from combined chunks" + ) + + +# --------------------------------------------------------------------------- +# MarkdownChunker +# --------------------------------------------------------------------------- + + +class TestMarkdownChunker: + def _chunker(self) -> MarkdownChunker: + return MarkdownChunker() + + @pytest.mark.asyncio + async def test_splits_by_headers(self): + """MarkdownChunker should produce one chunk per top-level section.""" + doc = Document(content=MARKDOWN_TEXT) + result = await self._chunker().chunk({}, [doc]) + + # Our sample has 4 sections (Introduction, Background, Details, Summary) + assert len(result[0].chunks) >= 3 + + @pytest.mark.asyncio + async def test_returns_chunk_instances(self): + """Every element in document.chunks must be a Chunk.""" + doc = Document(content=MARKDOWN_TEXT) + result = await self._chunker().chunk({}, [doc]) + + for chunk in result[0].chunks: + assert isinstance(chunk, Chunk) + + @pytest.mark.asyncio + async def test_chunk_content_nonempty(self): + """No chunk produced by the MarkdownChunker should be empty.""" + doc = Document(content=MARKDOWN_TEXT) + result = await self._chunker().chunk({}, [doc]) + + for chunk in result[0].chunks: + assert chunk.content.strip() != "" + + @pytest.mark.asyncio + async def test_chunk_ids_sequential(self): + """chunk_id values must form a sequence starting at 0.""" + doc = Document(content=MARKDOWN_TEXT) + result = await self._chunker().chunk({}, [doc]) + ids = [c.chunk_id for c in result[0].chunks] + assert ids == list(range(len(ids))) + + @pytest.mark.asyncio + async def test_headers_prepended_to_chunk_content(self): + """Fix for PR #323: header text must appear inside the chunk so retrieval + is context-aware even when the section body alone is ambiguous.""" + doc = Document(content=MARKDOWN_TEXT) + result = await self._chunker().chunk({}, [doc]) + + # The 'Background' section chunk must contain 'Background' in its content + background_chunks = [ + c for c in result[0].chunks if "Background" in c.content + ] + assert len(background_chunks) >= 1, ( + "Expected at least one chunk whose content includes the 'Background' header" + ) + + @pytest.mark.asyncio + async def test_nested_headers_included_in_chunk(self): + """A subsection chunk should include its ancestor header names (PR #323).""" + doc = Document(content=MARKDOWN_TEXT) + result = await self._chunker().chunk({}, [doc]) + + details_chunks = [c for c in result[0].chunks if "Details" in c.content] + assert len(details_chunks) >= 1, ( + "Chunk for the '### Details' subsection must include the header text" + ) + + @pytest.mark.asyncio + async def test_content_without_overlap_equals_content(self): + """MarkdownChunker has no overlap concept; both fields should be equal.""" + doc = Document(content=MARKDOWN_TEXT) + result = await self._chunker().chunk({}, [doc]) + + for chunk in result[0].chunks: + assert chunk.content == chunk.content_without_overlap + + @pytest.mark.asyncio + async def test_plain_text_no_headers_yields_one_chunk(self): + """Markdown text with no headers should produce a single chunk.""" + doc = Document(content="Just a plain paragraph without any headings.") + result = await self._chunker().chunk({}, [doc]) + + assert len(result[0].chunks) == 1 + + @pytest.mark.asyncio + async def test_skips_already_chunked_documents(self): + """Pre-chunked documents must not be re-chunked.""" + doc = Document(content=MARKDOWN_TEXT) + sentinel = Chunk( + content="sentinel", + chunk_id=0, + start_i=0, + end_i=8, + content_without_overlap="sentinel", + ) + doc.chunks.append(sentinel) + + result = await self._chunker().chunk({}, [doc]) + + assert len(result[0].chunks) == 1 + assert result[0].chunks[0].content == "sentinel" + + @pytest.mark.asyncio + async def test_multiple_documents_chunked_independently(self): + """Each document in the list must be chunked independently.""" + doc1 = Document(content=MARKDOWN_TEXT) + doc2 = Document(content="# Solo\n\nOnly one section here.") + + result = await self._chunker().chunk({}, [doc1, doc2]) + + assert len(result[0].chunks) >= 3, "First doc should have several sections" + assert len(result[1].chunks) == 1, "Second doc has one section" + + +# --------------------------------------------------------------------------- +# Cross-chunker contract tests +# --------------------------------------------------------------------------- + + +class TestChunkerContracts: + """Every chunker must satisfy a shared contract.""" + + CHUNKERS = [ + (TokenChunker, _make_token_config(50, 0)), + (SentenceChunker, _make_sentence_config(2, 0)), + (MarkdownChunker, {}), + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("chunker_cls,config", CHUNKERS) + async def test_returns_list_of_documents(self, chunker_cls, config): + doc = Document(content=MULTI_SENTENCE) + result = await chunker_cls().chunk(config, [doc]) + assert isinstance(result, list) + assert all(isinstance(d, Document) for d in result) + + @pytest.mark.asyncio + @pytest.mark.parametrize("chunker_cls,config", CHUNKERS) + async def test_chunks_are_chunk_instances(self, chunker_cls, config): + doc = Document(content=MULTI_SENTENCE) + result = await chunker_cls().chunk(config, [doc]) + for chunk in result[0].chunks: + assert isinstance(chunk, Chunk) + + @pytest.mark.asyncio + @pytest.mark.parametrize("chunker_cls,config", CHUNKERS) + async def test_no_empty_chunks(self, chunker_cls, config): + doc = Document(content=MULTI_SENTENCE) + result = await chunker_cls().chunk(config, [doc]) + for chunk in result[0].chunks: + assert chunk.content.strip() != "" diff --git a/goldenverba/tests/components/reader/__init__.py b/goldenverba/tests/components/reader/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/goldenverba/tests/components/reader/test_readers.py b/goldenverba/tests/components/reader/test_readers.py new file mode 100644 index 0000000000..9aa3ef3363 --- /dev/null +++ b/goldenverba/tests/components/reader/test_readers.py @@ -0,0 +1,740 @@ +""" +test_readers.py +=============== +Unit and integration tests for the reader system. + +Unit tests (Groups 1–7) run with no API keys and no network — all external +calls are mocked. Integration tests (Group 8) are opt-in: each test skips +unless the corresponding env var is set. + +How to enable integration tests +-------------------------------- +Set any combination of these in goldenverba/.env or your shell: + + UNSTRUCTURED_API_KEY=... + GITHUB_TOKEN=... + GITLAB_TOKEN=... + +HTMLReader integration test uses https://example.com — no key required, but +set READER_NETWORK_TESTS=1 to opt in (avoids surprises in air-gapped CI). +""" + +import asyncio +import base64 +import io +import json +import os + +import pytest +import pytest_asyncio +from unittest.mock import AsyncMock, MagicMock, patch, PropertyMock + +from goldenverba.components.reader.reader_manager import ReaderManager +from goldenverba.components.reader.BasicReader import BasicReader +from goldenverba.components.reader.HTMLReader import HTMLReader +from goldenverba.components.reader.GitReader import GitReader +from goldenverba.components.reader.UnstructuredAPI import UnstructuredReader +from goldenverba.components.reader.WhisperReader import WhisperReader +from goldenverba.server.types import FileConfig, FileStatus, RAGComponentClass, RAGComponentConfig +from goldenverba.components.types import InputConfig + + +# --------------------------------------------------------------------------- +# Integration test skip markers +# --------------------------------------------------------------------------- + +requires_unstructured = pytest.mark.skipif( + not os.getenv("UNSTRUCTURED_API_KEY"), + reason="Set UNSTRUCTURED_API_KEY to run Unstructured integration tests", +) +requires_github = pytest.mark.skipif( + not os.getenv("GITHUB_TOKEN"), + reason="Set GITHUB_TOKEN to run GitHub reader integration tests", +) +requires_gitlab = pytest.mark.skipif( + not os.getenv("GITLAB_TOKEN"), + reason="Set GITLAB_TOKEN to run GitLab reader integration tests", +) +requires_network = pytest.mark.skipif( + not os.getenv("READER_NETWORK_TESTS"), + reason="Set READER_NETWORK_TESTS=1 to run network-dependent reader tests", +) + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _b64(text: str) -> str: + return base64.b64encode(text.encode()).decode() + + +def _minimal_rag_class(name="Default") -> RAGComponentClass: + """Build the smallest valid RAGComponentClass Pydantic model.""" + component = RAGComponentConfig( + name=name, + variables=[], + library=[], + description="", + config={}, + type="", + available=True, + ) + return RAGComponentClass(selected=name, components={name: component}) + + +def _make_file_config( + filename="test.txt", + extension="txt", + content="hello world", + is_url=False, + overwrite=False, + rag_config=None, +): + return FileConfig( + fileID="file-001", + filename=filename, + isURL=is_url, + overwrite=overwrite, + extension=extension, + source="", + content=_b64(content) if extension != "" else content, + labels=[], + rag_config=rag_config or {"Reader": _minimal_rag_class()}, + file_size=len(content), + status=FileStatus.READY, + metadata="", + status_report={}, + ) + + +def _make_config(**overrides) -> dict: + """Return a minimal reader config dict.""" + return overrides + + +def _mock_aiohttp_response(status=200, json_data=None, text_data="<html>hi</html>"): + """Return a context-manager-compatible mock aiohttp response.""" + response = MagicMock() + response.status = status + response.raise_for_status = MagicMock() + response.json = AsyncMock(return_value=json_data or {}) + response.text = AsyncMock(return_value=text_data) + response.read = AsyncMock(return_value=text_data.encode() if isinstance(text_data, str) else text_data) + + if status >= 400: + from aiohttp import ClientResponseError + response.raise_for_status.side_effect = ClientResponseError( + request_info=MagicMock(), history=(), status=status + ) + + cm = MagicMock() + cm.__aenter__ = AsyncMock(return_value=response) + cm.__aexit__ = AsyncMock(return_value=False) + return cm, response + + +def _mock_session(get_response=None, post_response=None): + """Return a mock aiohttp.ClientSession context manager.""" + session = MagicMock() + if get_response is not None: + session.get = MagicMock(return_value=get_response) + if post_response is not None: + session.post = MagicMock(return_value=post_response) + + session_cm = MagicMock() + session_cm.__aenter__ = AsyncMock(return_value=session) + session_cm.__aexit__ = AsyncMock(return_value=False) + return session_cm, session + + +# --------------------------------------------------------------------------- +# 1. ReaderManager +# --------------------------------------------------------------------------- + + +class TestReaderManager: + def test_all_readers_registered(self): + manager = ReaderManager() + expected = {"Default", "HTML", "Git", "Unstructured IO", "Whisper"} + assert expected == set(manager.readers.keys()) + + @pytest.mark.asyncio + async def test_unknown_reader_raises(self): + manager = ReaderManager() + logger = MagicMock() + logger.send_report = AsyncMock() + fc = _make_file_config() + + with pytest.raises(Exception, match="Reader .* not found"): + await manager.load("NonExistentReader", fc, logger) + + @pytest.mark.asyncio + async def test_load_delegates_to_reader(self): + manager = ReaderManager() + logger = MagicMock() + logger.send_report = AsyncMock() + + fake_doc = MagicMock() + manager.readers["Default"] = MagicMock() + manager.readers["Default"].load = AsyncMock(return_value=[fake_doc]) + + fc = _make_file_config() + result = await manager.load("Default", fc, logger) + + assert result == [fake_doc] + manager.readers["Default"].load.assert_awaited_once() + + @pytest.mark.asyncio + async def test_load_sends_status_report(self): + manager = ReaderManager() + logger = MagicMock() + logger.send_report = AsyncMock() + + fake_doc = MagicMock() + manager.readers["Default"] = MagicMock() + manager.readers["Default"].load = AsyncMock(return_value=[fake_doc]) + + fc = _make_file_config() + await manager.load("Default", fc, logger) + + logger.send_report.assert_awaited() + + +# --------------------------------------------------------------------------- +# 2. BasicReader +# --------------------------------------------------------------------------- + + +class TestBasicReader: + @pytest.fixture + def reader(self): + return BasicReader() + + # --- text / code --- + + @pytest.mark.asyncio + async def test_loads_txt_file(self, reader): + fc = _make_file_config(extension="txt", content="hello world") + docs = await reader.load({}, fc) + assert len(docs) == 1 + assert "hello world" in docs[0].content + + @pytest.mark.asyncio + async def test_loads_py_file(self, reader): + fc = _make_file_config(filename="script.py", extension="py", content="print('hi')") + docs = await reader.load({}, fc) + assert len(docs) == 1 + assert "print('hi')" in docs[0].content + + @pytest.mark.asyncio + async def test_unknown_extension_falls_back_to_text(self, reader): + fc = _make_file_config(filename="foo.xyz", extension="xyz", content="raw bytes") + docs = await reader.load({}, fc) + assert len(docs) == 1 + + # --- PDF --- + + @pytest.mark.asyncio + async def test_load_pdf_returns_text(self, reader): + page = MagicMock() + page.extract_text.return_value = "Page content" + + with patch("goldenverba.components.reader.BasicReader.PdfReader") as MockPdf: + MockPdf.return_value.pages = [page] + result = await reader.load_pdf_file(b"fake-pdf-bytes") + + assert "Page content" in result + + @pytest.mark.asyncio + async def test_load_pdf_skips_none_pages(self, reader): + page_good = MagicMock() + page_good.extract_text.return_value = "Good page" + page_none = MagicMock() + page_none.extract_text.return_value = None + + with patch("goldenverba.components.reader.BasicReader.PdfReader") as MockPdf: + MockPdf.return_value.pages = [page_none, page_good] + result = await reader.load_pdf_file(b"fake-pdf-bytes") + + assert "None" not in result + assert "Good page" in result + + @pytest.mark.asyncio + async def test_load_pdf_raises_when_pypdf_missing(self, reader): + with patch("goldenverba.components.reader.BasicReader.PdfReader", None): + with pytest.raises(ImportError): + await reader.load_pdf_file(b"data") + + # --- DOCX --- + + @pytest.mark.asyncio + async def test_load_docx_returns_paragraph_text(self, reader): + para1, para2 = MagicMock(), MagicMock() + para1.text = "First paragraph" + para2.text = "Second paragraph" + + with patch("goldenverba.components.reader.BasicReader.docx") as mock_docx: + mock_docx.Document.return_value.paragraphs = [para1, para2] + result = await reader.load_docx_file(b"fake-docx") + + assert "First paragraph" in result + assert "Second paragraph" in result + + @pytest.mark.asyncio + async def test_load_docx_raises_when_docx_missing(self, reader): + with patch("goldenverba.components.reader.BasicReader.docx", None): + with pytest.raises(ImportError): + await reader.load_docx_file(b"data") + + # --- CSV --- + + @pytest.mark.asyncio + async def test_load_csv_formats_headers_and_rows(self, reader): + csv_bytes = b"name,age\nAlice,30\nBob,25" + result = await reader.load_csv_file(csv_bytes) + assert "Headers:" in result + assert "name" in result + assert "Alice" in result + + @pytest.mark.asyncio + async def test_load_csv_handles_empty_file(self, reader): + result = await reader.load_csv_file(b"") + assert "Empty" in result + + @pytest.mark.asyncio + async def test_load_csv_handles_mismatched_columns(self, reader): + # Row 2 has fewer columns than header — should not crash + csv_bytes = b"a,b,c\n1,2\n3,4,5" + result = await reader.load_csv_file(csv_bytes) + assert result # just needs to not raise + + # --- JSON --- + + @pytest.mark.asyncio + async def test_load_json_raises_on_invalid_json(self, reader): + fc = _make_file_config(extension="json") + bad_json = b"not valid json{" + with pytest.raises(ValueError, match="Invalid JSON"): + await reader.load_json_file(bad_json, fc) + + @pytest.mark.asyncio + async def test_load_json_falls_back_to_pretty_print(self, reader): + fc = _make_file_config(extension="json") + valid_json = json.dumps({"key": "value"}).encode() + + with patch("goldenverba.components.document.Document.from_json", return_value=None): + docs = await reader.load_json_file(valid_json, fc) + + assert len(docs) == 1 + assert "key" in docs[0].content + + +# --------------------------------------------------------------------------- +# 3. HTMLReader +# --------------------------------------------------------------------------- + + +class TestHTMLReader: + @pytest.fixture + def reader(self): + return HTMLReader() + + # --- extract_links --- + + def test_extract_links_same_domain(self, reader): + html = '<a href="/page">link</a><a href="https://example.com/other">other</a>' + links = reader.extract_links(html, "https://example.com/base") + assert all("example.com" in l for l in links) + + def test_extract_links_excludes_off_domain(self, reader): + html = '<a href="https://evil.com/steal">bad</a>' + links = reader.extract_links(html, "https://example.com/") + assert links == [] + + def test_extract_links_resolves_relative(self, reader): + html = '<a href="/about">about</a>' + links = reader.extract_links(html, "https://example.com/") + assert "https://example.com/about" in links + + # --- fetch_html_and_convert --- + + @pytest.mark.asyncio + async def test_fetch_returns_base64_html(self, reader): + html_text = "<h1>Hello</h1>" + get_cm, _ = _mock_aiohttp_response(text_data=html_text) + session = MagicMock() + session.get = MagicMock(return_value=get_cm) + + content_b64, size, raw = await reader.fetch_html_and_convert(session, "https://example.com", False) + + decoded = base64.b64decode(content_b64).decode() + assert decoded == html_text + assert size == len(html_text.encode()) + + @pytest.mark.asyncio + async def test_fetch_converts_to_markdown(self, reader): + html_text = "<h1>Title</h1>" + get_cm, _ = _mock_aiohttp_response(text_data=html_text) + session = MagicMock() + session.get = MagicMock(return_value=get_cm) + + with patch("goldenverba.components.reader.HTMLReader.md", return_value="# Title\n") as mock_md: + content_b64, _, _ = await reader.fetch_html_and_convert(session, "https://example.com", True) + mock_md.assert_called_once_with(html_text) + + decoded = base64.b64decode(content_b64).decode() + assert "Title" in decoded + + # --- load --- + + @pytest.mark.asyncio + async def test_load_returns_document_per_url(self, reader): + html_text = "<p>content</p>" + get_cm, _ = _mock_aiohttp_response(text_data=html_text) + session_cm, session = _mock_session(get_response=get_cm) + session.get = MagicMock(return_value=get_cm) + + config = { + "URLs": InputConfig(type="multi", value="", description="", values=["https://example.com"]), + "Convert To Markdown": InputConfig(type="bool", value=False, description="", values=[]), + "Recursive": InputConfig(type="bool", value=False, description="", values=[]), + "Max Depth": InputConfig(type="number", value=1, description="", values=[]), + } + + with patch("goldenverba.components.reader.HTMLReader.aiohttp.ClientSession", return_value=session_cm): + docs = await reader.load(config, _make_file_config()) + + assert len(docs) == 1 + + @pytest.mark.asyncio + async def test_load_continues_on_url_failure(self, reader): + """A failing URL should be skipped, not crash the whole load.""" + config = { + "URLs": InputConfig(type="multi", value="", description="", values=["https://bad-url.example"]), + "Convert To Markdown": InputConfig(type="bool", value=False, description="", values=[]), + "Recursive": InputConfig(type="bool", value=False, description="", values=[]), + "Max Depth": InputConfig(type="number", value=1, description="", values=[]), + } + + # Simulate a session whose get() always raises + session = MagicMock() + session.get = MagicMock(side_effect=Exception("network error")) + session_cm = MagicMock() + session_cm.__aenter__ = AsyncMock(return_value=session) + session_cm.__aexit__ = AsyncMock(return_value=False) + + with patch("goldenverba.components.reader.HTMLReader.aiohttp.ClientSession", return_value=session_cm): + docs = await reader.load(config, _make_file_config()) + + # No docs returned but no exception raised + assert docs == [] + + +# --------------------------------------------------------------------------- +# 4. GitReader +# --------------------------------------------------------------------------- + + +class TestGitReader: + @pytest.fixture + def reader(self): + return GitReader() + + def test_get_headers_github(self, reader): + headers = reader.get_headers("mytoken", "GitHub") + assert headers["Authorization"] == "token mytoken" + assert "Accept" in headers + + def test_get_headers_gitlab(self, reader): + headers = reader.get_headers("mytoken", "GitLab") + assert headers["Authorization"] == "Bearer mytoken" + + def test_get_token_reads_github_env(self, reader): + config = {} # no "Git Token" key + with patch.dict(os.environ, {"GITHUB_TOKEN": "gh-tok"}): + token = reader.get_token(config, "GitHub") + assert token == "gh-tok" + + def test_get_token_reads_gitlab_env(self, reader): + config = {} + with patch.dict(os.environ, {"GITLAB_TOKEN": "gl-tok"}): + token = reader.get_token(config, "GitLab") + assert token == "gl-tok" + + @pytest.mark.asyncio + async def test_fetch_docs_github_filters_by_extension_and_path(self, reader): + api_tree = { + "tree": [ + {"path": "src/main.py", "type": "blob"}, + {"path": "src/readme.md", "type": "blob"}, + {"path": "other/main.py", "type": "blob"}, + {"path": "src/image.png", "type": "blob"}, + ] + } + get_cm, _ = _mock_aiohttp_response(json_data=api_tree) + session_cm, session = _mock_session(get_response=get_cm) + session.get = MagicMock(return_value=get_cm) + + mock_reader = MagicMock() + mock_reader.extension = [".py", ".md"] + + with patch("goldenverba.components.reader.GitReader.aiohttp.ClientSession", return_value=session_cm): + paths = await reader.fetch_docs_github("https://api.github.com/...", "src", "token", mock_reader) + + # Should include src/*.py and src/*.md but not other/ or .png + assert "src/main.py" in paths + assert "src/readme.md" in paths + assert "other/main.py" not in paths + assert "src/image.png" not in paths + + @pytest.mark.asyncio + async def test_fetch_docs_gitlab_filters_blobs_by_extension(self, reader): + api_data = [ + {"path": "app.py", "type": "blob"}, + {"path": "app.png", "type": "blob"}, + {"path": "subdir", "type": "tree"}, + ] + get_cm, _ = _mock_aiohttp_response(json_data=api_data) + session_cm, session = _mock_session(get_response=get_cm) + session.get = MagicMock(return_value=get_cm) + + mock_reader = MagicMock() + mock_reader.extension = [".py"] + + with patch("goldenverba.components.reader.GitReader.aiohttp.ClientSession", return_value=session_cm): + paths = await reader.fetch_docs_gitlab("https://gitlab.com/...", "token", mock_reader) + + assert "app.py" in paths + assert "app.png" not in paths + assert "subdir" not in paths + + @pytest.mark.asyncio + async def test_download_file_github_returns_tuple(self, reader): + api_data = { + "content": _b64("print('hello')"), + "html_url": "https://github.com/owner/repo/blob/main/src/main.py", + "size": 15, + } + get_cm, _ = _mock_aiohttp_response(json_data=api_data) + session_cm, session = _mock_session(get_response=get_cm) + session.get = MagicMock(return_value=get_cm) + + with patch("goldenverba.components.reader.GitReader.aiohttp.ClientSession", return_value=session_cm): + content, link, size, ext = await reader.download_file_github( + "owner", "repo", "src/main.py", "main", "token" + ) + + assert content == _b64("print('hello')") + assert "github.com" in link + assert size == 15 + assert ext == "py" + + @pytest.mark.asyncio + async def test_download_file_gitlab_raises_on_error(self, reader): + get_cm, mock_resp = _mock_aiohttp_response(status=404, text_data="Not Found") + mock_resp.status = 404 + session_cm, session = _mock_session(get_response=get_cm) + session.get = MagicMock(return_value=get_cm) + + with patch("goldenverba.components.reader.GitReader.aiohttp.ClientSession", return_value=session_cm): + with pytest.raises(Exception, match="Failed to download"): + await reader.download_file_gitlab("owner", "repo", "missing.py", "main", "token") + + +# --------------------------------------------------------------------------- +# 5. UnstructuredReader +# --------------------------------------------------------------------------- + + +class TestUnstructuredReader: + @pytest.fixture + def reader(self): + return UnstructuredReader() + + def _config(self, strategy="auto", api_key="test-key", api_url="https://api.unstructuredapp.io/general/v0/general"): + return { + "Strategy": InputConfig(type="dropdown", value=strategy, description="", values=["auto", "hi_res", "ocr_only", "fast"]), + "API Key": InputConfig(type="password", value=api_key, description="", values=[]), + "API URL": InputConfig(type="text", value=api_url, description="", values=[]), + } + + @pytest.mark.asyncio + async def test_raises_on_invalid_strategy(self, reader): + fc = _make_file_config() + with pytest.raises(ValueError, match="Invalid strategy"): + await reader.load(self._config(strategy="invalid"), fc) + + @pytest.mark.asyncio + async def test_joins_chunk_texts(self, reader): + api_response = [{"text": "Hello "}, {"text": "world"}] + post_cm, _ = _mock_aiohttp_response(json_data=api_response) + session_cm, session = _mock_session(post_response=post_cm) + session.post = MagicMock(return_value=post_cm) + + with patch("goldenverba.components.reader.UnstructuredAPI.aiohttp.ClientSession", return_value=session_cm): + docs = await reader.load(self._config(), _make_file_config()) + + assert len(docs) == 1 + assert docs[0].content == "Hello world" + + @pytest.mark.asyncio + async def test_raises_on_api_error_detail(self, reader): + api_response = {"detail": "Invalid API key"} + post_cm, _ = _mock_aiohttp_response(json_data=api_response) + session_cm, session = _mock_session(post_response=post_cm) + session.post = MagicMock(return_value=post_cm) + + with patch("goldenverba.components.reader.UnstructuredAPI.aiohttp.ClientSession", return_value=session_cm): + with pytest.raises(Exception, match="API error"): + await reader.load(self._config(), _make_file_config()) + + +# --------------------------------------------------------------------------- +# 6. WhisperReader +# --------------------------------------------------------------------------- + + +class TestWhisperReader: + @pytest.fixture + def reader(self): + return WhisperReader() + + def _config(self, model_size="base", device="cpu"): + return { + "Model Size": InputConfig(type="dropdown", value=model_size, description="", values=["tiny", "base", "small", "medium", "large-v3"]), + "Device": InputConfig(type="dropdown", value=device, description="", values=["cpu", "cuda", "auto"]), + } + + def test_reader_name(self, reader): + assert reader.name == "Whisper" + + def test_no_env_required(self, reader): + assert reader.requires_env == [] + + def test_audio_extensions_present(self, reader): + for ext in [".mp3", ".wav", ".mp4", ".flac", ".ogg"]: + assert ext in reader.extension + + @pytest.mark.asyncio + async def test_raises_when_faster_whisper_missing(self, reader): + with patch("goldenverba.components.reader.WhisperReader.WhisperModel", None): + with pytest.raises(ImportError, match="faster-whisper"): + await reader.load(self._config(), _make_file_config(extension="wav", content="fake")) + + @pytest.mark.asyncio + async def test_uses_asyncio_to_thread(self, reader): + """_transcribe() must run in a thread, not block the event loop.""" + mock_segment = MagicMock() + mock_segment.text = "Hello world" + + with patch("goldenverba.components.reader.WhisperReader.WhisperModel"): + with patch("goldenverba.components.reader.WhisperReader.asyncio.to_thread", new_callable=AsyncMock) as mock_thread: + mock_thread.return_value = ([mock_segment], MagicMock()) + fc = _make_file_config(extension="wav", content="fake-audio") + docs = await reader.load(self._config(), fc) + + mock_thread.assert_awaited_once() + assert "Hello world" in docs[0].content + + @pytest.mark.asyncio + async def test_raises_on_empty_transcript(self, reader): + mock_segment = MagicMock() + mock_segment.text = " " # whitespace only → empty after strip + + with patch("goldenverba.components.reader.WhisperReader.WhisperModel"): + with patch("goldenverba.components.reader.WhisperReader.asyncio.to_thread", new_callable=AsyncMock) as mock_thread: + mock_thread.return_value = ([mock_segment], MagicMock()) + with pytest.raises(Exception, match="empty transcript"): + await reader.load(self._config(), _make_file_config(extension="wav", content="fake")) + + @pytest.mark.asyncio + async def test_cleans_up_temp_file_on_success(self, reader): + mock_segment = MagicMock() + mock_segment.text = "Some speech" + + with patch("goldenverba.components.reader.WhisperReader.WhisperModel"): + with patch("goldenverba.components.reader.WhisperReader.asyncio.to_thread", new_callable=AsyncMock) as mock_thread: + with patch("goldenverba.components.reader.WhisperReader.os.unlink") as mock_unlink: + mock_thread.return_value = ([mock_segment], MagicMock()) + await reader.load(self._config(), _make_file_config(extension="wav", content="fake")) + mock_unlink.assert_called_once() + + @pytest.mark.asyncio + async def test_cleans_up_temp_file_on_failure(self, reader): + with patch("goldenverba.components.reader.WhisperReader.WhisperModel"): + with patch("goldenverba.components.reader.WhisperReader.asyncio.to_thread", new_callable=AsyncMock) as mock_thread: + with patch("goldenverba.components.reader.WhisperReader.os.unlink") as mock_unlink: + mock_thread.side_effect = Exception("model crash") + with pytest.raises(Exception): + await reader.load(self._config(), _make_file_config(extension="wav", content="fake")) + mock_unlink.assert_called_once() + + +# --------------------------------------------------------------------------- +# 7. Integration tests (opt-in) +# --------------------------------------------------------------------------- + + +class TestHTMLReaderIntegration: + @requires_network + @pytest.mark.asyncio + async def test_loads_example_com(self): + reader = HTMLReader() + config = { + "URLs": InputConfig(type="multi", value="", description="", values=["https://example.com"]), + "Convert To Markdown": InputConfig(type="bool", value=False, description="", values=[]), + "Recursive": InputConfig(type="bool", value=False, description="", values=[]), + "Max Depth": InputConfig(type="number", value=1, description="", values=[]), + } + docs = await reader.load(config, _make_file_config()) + assert len(docs) == 1 + assert len(docs[0].content) > 0 + + +class TestUnstructuredIntegration: + @requires_unstructured + @pytest.mark.asyncio + async def test_loads_plain_text(self): + reader = UnstructuredReader() + config = { + "Strategy": InputConfig(type="dropdown", value="auto", description="", values=["auto", "hi_res", "ocr_only", "fast"]), + "API Key": InputConfig(type="password", value=os.environ["UNSTRUCTURED_API_KEY"], description="", values=[]), + # Always use the current production URL regardless of any stale env var + "API URL": InputConfig( + type="text", + value="https://api.unstructuredapp.io/general/v0/general", + description="", + values=[], + ), + } + fc = _make_file_config(filename="test.txt", extension="txt", content="Hello from Verba integration test.") + try: + docs = await reader.load(config, fc) + except Exception as e: + msg = str(e) + if "Cannot connect to host" in msg or "nodename nor servname" in msg: + pytest.skip(f"Unstructured API unreachable: {e}") + if "401" in msg or "Unauthorized" in msg: + pytest.skip(f"Unstructured API key invalid or expired: {e}") + raise + assert len(docs) == 1 + assert len(docs[0].content) > 0 + + +class TestGitReaderIntegration: + @requires_github + @pytest.mark.asyncio + async def test_reads_public_github_repo(self): + reader = GitReader() + config = { + "Platform": InputConfig(type="dropdown", value="GitHub", description="", values=["GitHub", "GitLab"]), + "Owner": InputConfig(type="text", value="weaviate", description="", values=[]), + "Name": InputConfig(type="text", value="Verba", description="", values=[]), + "Branch": InputConfig(type="text", value="main", description="", values=[]), + "Path": InputConfig(type="text", value="README.md", description="", values=[]), + "Git Token": InputConfig(type="password", value=os.environ["GITHUB_TOKEN"], description="", values=[]), + } + docs = await reader.load(config, _make_file_config()) + assert len(docs) >= 1 diff --git a/goldenverba/tests/chunk/test_chunk.py b/goldenverba/tests/components/test_chunk.py similarity index 100% rename from goldenverba/tests/chunk/test_chunk.py rename to goldenverba/tests/components/test_chunk.py diff --git a/goldenverba/tests/components/test_client_manager.py b/goldenverba/tests/components/test_client_manager.py new file mode 100644 index 0000000000..79f0650fd7 --- /dev/null +++ b/goldenverba/tests/components/test_client_manager.py @@ -0,0 +1,274 @@ +""" +test_client_manager.py +====================== +Unit tests for ClientManager. + +Strategy: instantiate a real ClientManager, then replace self.manager with a +MagicMock so no real Weaviate connections are opened. WeaviateAsyncClient +instances in the pool are also MagicMocks with AsyncMock.is_ready(). + +Groups: + 1. hash_credentials — deterministic, credential-isolated hashing + 2. get_or_create_lock — idempotent lock creation + 3. connect — cache hit, cache miss, env-var fallback + 4. disconnect — closes all pooled clients + 5. clean_up — evicts stale-by-time and unhealthy clients +""" + +import asyncio +import os +from datetime import datetime, timedelta +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from goldenverba.components.client_manager import ClientManager +from goldenverba.server.types import Credentials + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _creds(deployment="Weaviate", url="https://example.weaviate.network", key="secret"): + return Credentials(deployment=deployment, url=url, key=key) + + +def _mock_weaviate_client(ready=True): + client = MagicMock() + client.is_ready = AsyncMock(return_value=ready) + return client + + +@pytest.fixture +def cm(): + """ClientManager with VerbaManager replaced by a MagicMock.""" + manager = ClientManager() + manager.manager = MagicMock() + return manager + + +# --------------------------------------------------------------------------- +# 1. hash_credentials +# --------------------------------------------------------------------------- + + +class TestHashCredentials: + def test_same_credentials_produce_same_hash(self, cm): + creds = _creds() + assert cm.hash_credentials(creds) == cm.hash_credentials(creds) + + def test_different_url_produces_different_hash(self, cm): + a = _creds(url="https://a.example.com") + b = _creds(url="https://b.example.com") + assert cm.hash_credentials(a) != cm.hash_credentials(b) + + def test_different_key_produces_different_hash(self, cm): + a = _creds(key="key-a") + b = _creds(key="key-b") + assert cm.hash_credentials(a) != cm.hash_credentials(b) + + def test_different_deployment_produces_different_hash(self, cm): + a = _creds(deployment="Weaviate") + b = _creds(deployment="Docker") + assert cm.hash_credentials(a) != cm.hash_credentials(b) + + def test_hash_is_64_char_hex(self, cm): + h = cm.hash_credentials(_creds()) + assert len(h) == 64 + assert all(c in "0123456789abcdef" for c in h) + + +# --------------------------------------------------------------------------- +# 2. get_or_create_lock +# --------------------------------------------------------------------------- + + +class TestGetOrCreateLock: + def test_returns_asyncio_lock(self, cm): + lock = cm.get_or_create_lock("abc") + assert isinstance(lock, asyncio.Lock) + + def test_same_key_returns_same_lock(self, cm): + lock1 = cm.get_or_create_lock("abc") + lock2 = cm.get_or_create_lock("abc") + assert lock1 is lock2 + + def test_different_keys_return_different_locks(self, cm): + lock1 = cm.get_or_create_lock("abc") + lock2 = cm.get_or_create_lock("xyz") + assert lock1 is not lock2 + + +# --------------------------------------------------------------------------- +# 3. connect +# --------------------------------------------------------------------------- + + +class TestConnect: + @pytest.mark.asyncio + async def test_cache_miss_opens_new_connection(self, cm): + fake_client = _mock_weaviate_client() + cm.manager.connect = AsyncMock(return_value=fake_client) + creds = _creds() + + result = await cm.connect(creds) + + assert result is fake_client + cm.manager.connect.assert_awaited_once() + + @pytest.mark.asyncio + async def test_cache_hit_returns_existing_client(self, cm): + fake_client = _mock_weaviate_client() + cm.manager.connect = AsyncMock(return_value=fake_client) + creds = _creds() + + # First call stores it; second call should reuse. + first = await cm.connect(creds) + second = await cm.connect(creds) + + assert first is second + # manager.connect only called once despite two cm.connect calls + cm.manager.connect.assert_awaited_once() + + @pytest.mark.asyncio + async def test_different_credentials_open_separate_connections(self, cm): + client_a = _mock_weaviate_client() + client_b = _mock_weaviate_client() + cm.manager.connect = AsyncMock(side_effect=[client_a, client_b]) + + result_a = await cm.connect(_creds(key="key-a")) + result_b = await cm.connect(_creds(key="key-b")) + + assert result_a is client_a + assert result_b is client_b + assert cm.manager.connect.await_count == 2 + + @pytest.mark.asyncio + async def test_env_var_fallback_when_empty_credentials(self, cm): + fake_client = _mock_weaviate_client() + cm.manager.connect = AsyncMock(return_value=fake_client) + empty_creds = _creds(url="", key="") + + with patch.dict(os.environ, { + "WEAVIATE_URL_VERBA": "https://env.weaviate.network", + "WEAVIATE_API_KEY_VERBA": "env-key", + }): + await cm.connect(empty_creds) + + # The credentials passed to manager.connect should use the env values. + called_creds = cm.manager.connect.call_args.args[0] + assert called_creds.url == "https://env.weaviate.network" + assert called_creds.key == "env-key" + + @pytest.mark.asyncio + async def test_connect_does_not_mutate_caller_credentials(self, cm): + cm.manager.connect = AsyncMock(return_value=_mock_weaviate_client()) + original_url = "" + creds = _creds(url=original_url, key="") + + with patch.dict(os.environ, {"WEAVIATE_URL_VERBA": "https://env.example.com"}): + await cm.connect(creds) + + # Caller's object should be unchanged. + assert creds.url == original_url + + +# --------------------------------------------------------------------------- +# 4. disconnect +# --------------------------------------------------------------------------- + + +class TestDisconnect: + @pytest.mark.asyncio + async def test_disconnects_all_clients(self, cm): + client_a = _mock_weaviate_client() + client_b = _mock_weaviate_client() + cm.manager.disconnect = AsyncMock() + cm.clients = { + "hash-a": {"client": client_a, "timestamp": datetime.now()}, + "hash-b": {"client": client_b, "timestamp": datetime.now()}, + } + + await cm.disconnect() + + assert cm.manager.disconnect.await_count == 2 + disconnected = {call.args[0] for call in cm.manager.disconnect.call_args_list} + assert client_a in disconnected + assert client_b in disconnected + + @pytest.mark.asyncio + async def test_disconnect_with_empty_pool_is_a_noop(self, cm): + cm.manager.disconnect = AsyncMock() + await cm.disconnect() + cm.manager.disconnect.assert_not_called() + + +# --------------------------------------------------------------------------- +# 5. clean_up +# --------------------------------------------------------------------------- + + +class TestCleanUp: + @pytest.mark.asyncio + async def test_removes_stale_client_by_time(self, cm): + stale_client = _mock_weaviate_client(ready=True) + cm.manager.disconnect = AsyncMock() + old_timestamp = datetime.now() - timedelta(minutes=cm.max_time + 1) + cm.clients = { + "stale": {"client": stale_client, "timestamp": old_timestamp}, + } + + await cm.clean_up() + + assert "stale" not in cm.clients + cm.manager.disconnect.assert_awaited_once_with(stale_client) + + @pytest.mark.asyncio + async def test_removes_unhealthy_client(self, cm): + unhealthy_client = _mock_weaviate_client(ready=False) + cm.manager.disconnect = AsyncMock() + cm.clients = { + "unhealthy": {"client": unhealthy_client, "timestamp": datetime.now()}, + } + + await cm.clean_up() + + assert "unhealthy" not in cm.clients + cm.manager.disconnect.assert_awaited_once_with(unhealthy_client) + + @pytest.mark.asyncio + async def test_keeps_fresh_healthy_client(self, cm): + healthy_client = _mock_weaviate_client(ready=True) + cm.manager.disconnect = AsyncMock() + cm.clients = { + "healthy": {"client": healthy_client, "timestamp": datetime.now()}, + } + + await cm.clean_up() + + assert "healthy" in cm.clients + cm.manager.disconnect.assert_not_called() + + @pytest.mark.asyncio + async def test_mixed_pool_only_removes_bad_clients(self, cm): + good_client = _mock_weaviate_client(ready=True) + bad_client = _mock_weaviate_client(ready=False) + cm.manager.disconnect = AsyncMock() + cm.clients = { + "good": {"client": good_client, "timestamp": datetime.now()}, + "bad": {"client": bad_client, "timestamp": datetime.now()}, + } + + await cm.clean_up() + + assert "good" in cm.clients + assert "bad" not in cm.clients + cm.manager.disconnect.assert_awaited_once_with(bad_client) + + @pytest.mark.asyncio + async def test_empty_pool_clean_up_is_a_noop(self, cm): + cm.manager.disconnect = AsyncMock() + await cm.clean_up() + cm.manager.disconnect.assert_not_called() diff --git a/goldenverba/tests/document/test_document.py b/goldenverba/tests/components/test_document.py similarity index 60% rename from goldenverba/tests/document/test_document.py rename to goldenverba/tests/components/test_document.py index d3ac6e92de..addca16d82 100644 --- a/goldenverba/tests/document/test_document.py +++ b/goldenverba/tests/components/test_document.py @@ -55,10 +55,68 @@ def test_document_json_serialization(): assert restored_doc.metadata == original_doc.metadata +def _make_file_config(**overrides) -> FileConfig: + """Return a minimal valid FileConfig for testing.""" + from goldenverba.server.types import FileStatus + + defaults = dict( + fileID="fc-001", + filename="sample.txt", + isURL=False, + overwrite=False, + extension=".txt", + source="local", + content="Hello world", + labels=["tag1"], + rag_config={}, + file_size=11, + status=FileStatus.READY, + metadata="some metadata", + status_report={}, + ) + defaults.update(overrides) + return FileConfig(**defaults) + + def test_create_document_from_file_config(): - """Test document creation from FileConfig""" - # TODO: Add test - assert True + """Test document creation from FileConfig maps all fields correctly.""" + fc = _make_file_config() + doc = create_document("Hello world", fc) + + assert doc.title == fc.filename + assert doc.content == "Hello world" + assert doc.extension == fc.extension + assert doc.labels == fc.labels + assert doc.source == fc.source + assert doc.fileSize == fc.file_size + assert doc.metadata == fc.metadata + assert doc.meta == {} + + +def test_create_document_empty_content(): + """create_document with empty content should produce a Document with empty content.""" + fc = _make_file_config(filename="empty.txt", file_size=0) + doc = create_document("", fc) + + assert doc.content == "" + assert doc.title == "empty.txt" + + +def test_create_document_preserves_labels(): + """Labels from FileConfig must be forwarded to the Document unchanged.""" + fc = _make_file_config(labels=["invoice", "2024", "finance"]) + doc = create_document("some content", fc) + + assert doc.labels == ["invoice", "2024", "finance"] + + +def test_create_document_url_source(): + """create_document should work when source is a URL string.""" + fc = _make_file_config(source="https://example.com/doc.txt", isURL=True) + doc = create_document("Content from URL", fc) + + assert doc.source == "https://example.com/doc.txt" + assert doc.content == "Content from URL" def test_document_with_large_content(): diff --git a/goldenverba/tests/components/test_verba_manager.py b/goldenverba/tests/components/test_verba_manager.py new file mode 100644 index 0000000000..c8d10d864a --- /dev/null +++ b/goldenverba/tests/components/test_verba_manager.py @@ -0,0 +1,628 @@ +""" +test_verba_manager.py +===================== +Unit tests for VerbaManager. + +Strategy: instantiate a real VerbaManager, then replace sub-managers on the +instance with MagicMock / AsyncMock objects. This avoids patching import paths +and keeps each test focused on the behaviour being asserted. + +Groups: + 1. Pure unit tests — no mocks (create_user_config, _keys_match, verify_config, + get_deployments, verify_installed_libraries, verify_variables) + 2. Config round-trips — mock weaviate_manager for load/set/reset config methods + 3. Import pipeline — mock reader, chunker, embedder, weaviate managers + 4. RAG pipeline — mock embedder, retriever, generator managers +""" + +import os +import pytest +from unittest.mock import MagicMock, AsyncMock, patch + +from goldenverba.components.verba_manager import VerbaManager +from goldenverba.server.types import ( + FileConfig, + FileStatus, + ChunkScore, + RAGComponentClass, +) + + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def manager(): + """Real VerbaManager with weaviate_manager replaced by a MagicMock.""" + m = VerbaManager() + m.weaviate_manager = MagicMock() + return m + + +@pytest.fixture +def mock_logger(): + """Minimal LoggerManager stand-in.""" + logger = MagicMock() + logger.send_report = AsyncMock() + return logger + + +def _minimal_config_setting(description="desc", values=None): + """Return a minimal config-setting dict that verify_config can walk.""" + return {"description": description, "values": values or ["a", "b"]} + + +def _make_config(*category_keys): + """ + Build a minimal RAG config dict for verify_config tests. + Each category has one component ("comp") with one setting ("key"). + """ + return { + cat: { + "components": { + "comp": { + "config": { + "key": _minimal_config_setting() + } + } + } + } + for cat in category_keys + } + + +def _make_rag_component_class(selected="MockEmbedder", model_value="text-embedding-3-small"): + """Return a RAGComponentClass-like MagicMock for use in rag_config dicts.""" + config_setting = MagicMock() + config_setting.value = model_value + + component = MagicMock() + component.config = {"Model": config_setting} + + rag_class = MagicMock(spec=RAGComponentClass) + rag_class.selected = selected + rag_class.components = {selected: component} + return rag_class + + +# --------------------------------------------------------------------------- +# 1. Pure unit tests +# --------------------------------------------------------------------------- + + +class TestCreateUserConfig: + def test_returns_default(self, manager): + assert manager.create_user_config() == {"getting_started": False} + + +class TestKeysMatch: + def test_matching_keys_returns_true(self, manager): + assert manager._keys_match({"a": 1, "b": 2}, {"a": 3, "b": 4}, "test") is True + + def test_mismatched_keys_returns_false(self, manager): + assert manager._keys_match({"a": 1}, {"b": 1}, "test") is False + + def test_empty_dicts_match(self, manager): + assert manager._keys_match({}, {}, "test") is True + + def test_subset_returns_false(self, manager): + assert manager._keys_match({"a": 1}, {"a": 1, "b": 2}, "test") is False + + +class TestVerifyConfig: + def test_matching_configs_returns_true(self, manager): + cfg = _make_config("Reader", "Chunker") + assert manager.verify_config(cfg, cfg) is True + + def test_category_mismatch_returns_false(self, manager): + a = _make_config("Reader") + b = _make_config("Chunker") + assert manager.verify_config(a, b) is False + + def test_component_mismatch_returns_false(self, manager): + base = _make_config("Reader") + drift = { + "Reader": { + "components": { + "other_comp": {"config": {"key": _minimal_config_setting()}} + } + } + } + assert manager.verify_config(base, drift) is False + + def test_config_key_mismatch_returns_false(self, manager): + a = _make_config("Reader") + b = { + "Reader": { + "components": { + "comp": { + "config": { + "different_key": _minimal_config_setting() + } + } + } + } + } + assert manager.verify_config(a, b) is False + + def test_description_mismatch_returns_false(self, manager): + a = _make_config("Reader") + b = { + "Reader": { + "components": { + "comp": { + "config": { + "key": _minimal_config_setting(description="changed") + } + } + } + } + } + assert manager.verify_config(a, b) is False + + def test_values_mismatch_returns_false(self, manager): + a = _make_config("Reader") + b = { + "Reader": { + "components": { + "comp": { + "config": { + "key": _minimal_config_setting(values=["x", "y"]) + } + } + } + } + } + assert manager.verify_config(a, b) is False + + def test_malformed_config_returns_false(self, manager): + assert manager.verify_config({"bad": "data"}, _make_config("Reader")) is False + + def test_demo_mode_always_true(self, manager): + a = _make_config("Reader") + b = _make_config("Chunker") # deliberately mismatched + with patch.dict(os.environ, {"VERBA_PRODUCTION": "Demo"}): + assert manager.verify_config(a, b) is True + + +class TestGetDeployments: + @pytest.mark.asyncio + async def test_returns_env_vars_when_set(self, manager): + with patch.dict(os.environ, { + "WEAVIATE_URL_VERBA": "https://my.weaviate.io", + "WEAVIATE_API_KEY_VERBA": "secret", + }): + result = await manager.get_deployments() + assert result["WEAVIATE_URL_VERBA"] == "https://my.weaviate.io" + assert result["WEAVIATE_API_KEY_VERBA"] == "secret" + + @pytest.mark.asyncio + async def test_returns_empty_strings_when_unset(self, manager): + env = {k: v for k, v in os.environ.items() + if k not in ("WEAVIATE_URL_VERBA", "WEAVIATE_API_KEY_VERBA")} + with patch.dict(os.environ, env, clear=True): + result = await manager.get_deployments() + assert result["WEAVIATE_URL_VERBA"] == "" + assert result["WEAVIATE_API_KEY_VERBA"] == "" + + +class TestVerifyInstalledLibraries: + def test_importable_library_marked_true(self, manager): + # Inject a fake component whose requires_library contains a stdlib module. + fake_component = MagicMock() + fake_component.requires_library = ["os"] + fake_component.requires_env = [] + manager.reader_manager.readers = {"fake": fake_component} + # Reset the other managers so they contribute nothing. + for attr in ("chunkers", "embedders", "retrievers", "generators"): + mgr_name = attr.replace("ers", "er_manager").replace("ors", "or_manager") + manager.chunker_manager.chunkers = {} + manager.embedder_manager.embedders = {} + manager.retriever_manager.retrievers = {} + manager.generator_manager.generators = {} + + manager.installed_libraries = {} + manager.verify_installed_libraries() + assert manager.installed_libraries.get("os") is True + + def test_missing_library_marked_false(self, manager): + fake_component = MagicMock() + fake_component.requires_library = ["__nonexistent_lib__"] + fake_component.requires_env = [] + manager.reader_manager.readers = {"fake": fake_component} + manager.chunker_manager.chunkers = {} + manager.embedder_manager.embedders = {} + manager.retriever_manager.retrievers = {} + manager.generator_manager.generators = {} + + manager.installed_libraries = {} + manager.verify_installed_libraries() + assert manager.installed_libraries.get("__nonexistent_lib__") is False + + +class TestVerifyVariables: + def test_set_env_var_marked_true(self, manager): + fake_component = MagicMock() + fake_component.requires_env = ["TEST_VAR_VERBA"] + fake_component.requires_library = [] + manager.reader_manager.readers = {"fake": fake_component} + manager.chunker_manager.chunkers = {} + manager.embedder_manager.embedders = {} + manager.retriever_manager.retrievers = {} + manager.generator_manager.generators = {} + + with patch.dict(os.environ, {"TEST_VAR_VERBA": "value"}): + manager.environment_variables = {} + manager.verify_variables() + assert manager.environment_variables.get("TEST_VAR_VERBA") is True + + def test_unset_env_var_marked_false(self, manager): + fake_component = MagicMock() + fake_component.requires_env = ["MISSING_VAR_VERBA_XYZ"] + fake_component.requires_library = [] + manager.reader_manager.readers = {"fake": fake_component} + manager.chunker_manager.chunkers = {} + manager.embedder_manager.embedders = {} + manager.retriever_manager.retrievers = {} + manager.generator_manager.generators = {} + + env = {k: v for k, v in os.environ.items() if k != "MISSING_VAR_VERBA_XYZ"} + with patch.dict(os.environ, env, clear=True): + manager.environment_variables = {} + manager.verify_variables() + assert manager.environment_variables.get("MISSING_VAR_VERBA_XYZ") is False + + +# --------------------------------------------------------------------------- +# 2. Config round-trip tests +# --------------------------------------------------------------------------- + + +class TestLoadRagConfig: + @pytest.mark.asyncio + async def test_returns_stored_config_when_valid(self, manager): + """Stored config passes verify_config → returned as-is; no write.""" + # create_config() returns a real config; mirror it as the "stored" one. + real_config = manager.create_config() + manager.weaviate_manager.get_config = AsyncMock(return_value=real_config) + manager.weaviate_manager.set_config = AsyncMock() + + result = await manager.load_rag_config(client=MagicMock()) + + assert result is real_config + manager.weaviate_manager.set_config.assert_not_called() + + @pytest.mark.asyncio + async def test_regenerates_config_on_schema_drift(self, manager): + """Stored config fails verify_config → fresh config written and returned.""" + stale_config = {"completely": "wrong"} + manager.weaviate_manager.get_config = AsyncMock(return_value=stale_config) + manager.weaviate_manager.set_config = AsyncMock() + + result = await manager.load_rag_config(client=MagicMock()) + + # set_config called with the RAG UUID + manager.weaviate_manager.set_config.assert_awaited_once() + call_args = manager.weaviate_manager.set_config.call_args + assert call_args.args[1] == manager.rag_config_uuid + # Returned config is a fresh one (has all 5 pipeline categories). + assert set(result.keys()) == {"Reader", "Chunker", "Embedder", "Retriever", "Generator"} + + @pytest.mark.asyncio + async def test_returns_fresh_config_when_none_stored(self, manager): + """No stored config → fresh config returned; nothing written.""" + manager.weaviate_manager.get_config = AsyncMock(return_value=None) + manager.weaviate_manager.set_config = AsyncMock() + + result = await manager.load_rag_config(client=MagicMock()) + + manager.weaviate_manager.set_config.assert_not_called() + assert set(result.keys()) == {"Reader", "Chunker", "Embedder", "Retriever", "Generator"} + + +class TestLoadThemeConfig: + @pytest.mark.asyncio + async def test_returns_theme_and_themes_when_stored(self, manager): + stored = {"theme": "dark", "themes": {"dark": {}, "light": {}}} + manager.weaviate_manager.get_config = AsyncMock(return_value=stored) + + theme, themes = await manager.load_theme_config(client=MagicMock()) + + assert theme == "dark" + assert themes == {"dark": {}, "light": {}} + + @pytest.mark.asyncio + async def test_returns_none_tuple_when_nothing_stored(self, manager): + manager.weaviate_manager.get_config = AsyncMock(return_value=None) + + theme, themes = await manager.load_theme_config(client=MagicMock()) + + assert theme is None + assert themes is None + + +class TestLoadUserConfig: + @pytest.mark.asyncio + async def test_returns_stored_config(self, manager): + stored = {"getting_started": True, "extra": "data"} + manager.weaviate_manager.get_config = AsyncMock(return_value=stored) + + result = await manager.load_user_config(client=MagicMock()) + + assert result == stored + + @pytest.mark.asyncio + async def test_returns_default_when_nothing_stored(self, manager): + manager.weaviate_manager.get_config = AsyncMock(return_value=None) + + result = await manager.load_user_config(client=MagicMock()) + + assert result == {"getting_started": False} + + +class TestSetConfigs: + @pytest.mark.asyncio + async def test_set_rag_config_uses_correct_uuid(self, manager): + manager.weaviate_manager.set_config = AsyncMock() + client = MagicMock() + payload = {"some": "config"} + + await manager.set_rag_config(client, payload) + + manager.weaviate_manager.set_config.assert_awaited_once_with( + client, manager.rag_config_uuid, payload + ) + + @pytest.mark.asyncio + async def test_set_theme_config_uses_correct_uuid(self, manager): + manager.weaviate_manager.set_config = AsyncMock() + client = MagicMock() + payload = {"theme": "light"} + + await manager.set_theme_config(client, payload) + + manager.weaviate_manager.set_config.assert_awaited_once_with( + client, manager.theme_config_uuid, payload + ) + + @pytest.mark.asyncio + async def test_set_user_config_uses_correct_uuid(self, manager): + manager.weaviate_manager.set_config = AsyncMock() + client = MagicMock() + payload = {"getting_started": True} + + await manager.set_user_config(client, payload) + + manager.weaviate_manager.set_config.assert_awaited_once_with( + client, manager.user_config_uuid, payload + ) + + +class TestResetConfigs: + @pytest.mark.asyncio + async def test_reset_rag_config_uses_correct_uuid(self, manager): + manager.weaviate_manager.reset_config = AsyncMock() + client = MagicMock() + + await manager.reset_rag_config(client) + + manager.weaviate_manager.reset_config.assert_awaited_once_with( + client, manager.rag_config_uuid + ) + + @pytest.mark.asyncio + async def test_reset_theme_config_uses_correct_uuid(self, manager): + manager.weaviate_manager.reset_config = AsyncMock() + client = MagicMock() + + await manager.reset_theme_config(client) + + manager.weaviate_manager.reset_config.assert_awaited_once_with( + client, manager.theme_config_uuid + ) + + @pytest.mark.asyncio + async def test_reset_user_config_uses_correct_uuid(self, manager): + manager.weaviate_manager.reset_config = AsyncMock() + client = MagicMock() + + await manager.reset_user_config(client) + + manager.weaviate_manager.reset_config.assert_awaited_once_with( + client, manager.user_config_uuid + ) + + +# --------------------------------------------------------------------------- +# 3. Import pipeline tests +# --------------------------------------------------------------------------- + + +def _make_file_config(filename="test.txt", overwrite=False): + """Minimal FileConfig for import pipeline tests.""" + reader_class = _make_rag_component_class(selected="Basic") + chunker_class = _make_rag_component_class(selected="Token") + embedder_class = _make_rag_component_class(selected="OpenAI") + + return FileConfig( + fileID="file-001", + filename=filename, + isURL=False, + overwrite=overwrite, + extension=".txt", + source="", + content="hello world", + labels=[], + rag_config={ + "Reader": reader_class, + "Chunker": chunker_class, + "Embedder": embedder_class, + }, + file_size=11, + status=FileStatus.READY, + metadata="", + status_report={}, + ) + + +class TestImportDocument: + @pytest.mark.asyncio + async def test_happy_path_single_document(self, manager, mock_logger): + """Reader → chunk → embed → store: logger receives STARTING then DONE.""" + from goldenverba.components.document import Document + from goldenverba.components.chunk import Chunk + + chunk = Chunk(content="hello world", chunk_id="0") + doc = Document(title="test.txt", content="hello world", extension=".txt") + doc.chunks = [chunk] + + # Mock reader + manager.reader_manager.load = AsyncMock(return_value=[doc]) + # Mock chunker + manager.chunker_manager.chunk = AsyncMock(return_value=[doc]) + # Mock embedder + manager.embedder_manager.vectorize = AsyncMock(return_value=[doc]) + # Mock weaviate + manager.weaviate_manager.exist_document_name = AsyncMock(return_value=None) + manager.weaviate_manager.import_document = AsyncMock() + + file_config = _make_file_config() + await manager.import_document(MagicMock(), file_config, mock_logger) + + # First status: STARTING + first_call = mock_logger.send_report.call_args_list[0] + assert first_call.kwargs["status"] == FileStatus.STARTING + # Last status: DONE (from process_single_document or import_document level) + statuses = [call.kwargs["status"] for call in mock_logger.send_report.call_args_list] + assert FileStatus.DONE in statuses + + @pytest.mark.asyncio + async def test_duplicate_no_overwrite_reports_error(self, manager, mock_logger): + """Existing document + overwrite=False → ERROR logged; delete NOT called.""" + manager.weaviate_manager.exist_document_name = AsyncMock(return_value="existing-uuid") + manager.weaviate_manager.delete_document = AsyncMock() + + file_config = _make_file_config(overwrite=False) + await manager.import_document(MagicMock(), file_config, mock_logger) + + statuses = [call.kwargs["status"] for call in mock_logger.send_report.call_args_list] + assert FileStatus.ERROR in statuses + manager.weaviate_manager.delete_document.assert_not_called() + + @pytest.mark.asyncio + async def test_duplicate_with_overwrite_deletes_then_imports(self, manager, mock_logger): + """Existing document + overwrite=True → delete called before import.""" + from goldenverba.components.document import Document + from goldenverba.components.chunk import Chunk + + chunk = Chunk(content="hello world", chunk_id="0") + doc = Document(title="test.txt", content="hello world", extension=".txt") + doc.chunks = [chunk] + + manager.weaviate_manager.exist_document_name = AsyncMock(return_value="existing-uuid") + manager.weaviate_manager.delete_document = AsyncMock() + manager.reader_manager.load = AsyncMock(return_value=[doc]) + manager.chunker_manager.chunk = AsyncMock(return_value=[doc]) + manager.embedder_manager.vectorize = AsyncMock(return_value=[doc]) + manager.weaviate_manager.import_document = AsyncMock() + + file_config = _make_file_config(overwrite=True) + await manager.import_document(MagicMock(), file_config, mock_logger) + + # Both import_document and process_single_document check for duplicates, + # so delete_document is called once per duplicate check (two total here). + assert manager.weaviate_manager.delete_document.await_count >= 1 + statuses = [call.kwargs["status"] for call in mock_logger.send_report.call_args_list] + assert FileStatus.ERROR not in statuses + + +# --------------------------------------------------------------------------- +# 4. RAG pipeline tests +# --------------------------------------------------------------------------- + + +class TestRetrieveChunks: + @pytest.mark.asyncio + async def test_returns_documents_and_context(self, manager): + dummy_vector = [0.1, 0.2, 0.3] + expected_docs = [{"uuid": "doc-1", "chunks": []}] + expected_context = "some retrieved context" + + manager.embedder_manager.vectorize_query = AsyncMock(return_value=dummy_vector) + manager.retriever_manager.retrieve = AsyncMock( + return_value=(expected_docs, expected_context) + ) + manager.weaviate_manager.add_suggestion = AsyncMock() + + rag_config = { + "Retriever": _make_rag_component_class(selected="Window"), + "Embedder": _make_rag_component_class(selected="OpenAI"), + } + + docs, context = await manager.retrieve_chunks( + client=MagicMock(), + query="what is Verba?", + rag_config=rag_config, + ) + + assert docs == expected_docs + assert context == expected_context + + @pytest.mark.asyncio + async def test_short_query_skips_suggestion(self, manager): + """Query shorter than 3 chars → add_suggestion not called.""" + manager.embedder_manager.vectorize_query = AsyncMock(return_value=[0.0]) + manager.retriever_manager.retrieve = AsyncMock(return_value=([], "")) + manager.weaviate_manager.add_suggestion = AsyncMock() + + rag_config = { + "Retriever": _make_rag_component_class(selected="Window"), + "Embedder": _make_rag_component_class(selected="OpenAI"), + } + + await manager.retrieve_chunks( + client=MagicMock(), + query="hi", + rag_config=rag_config, + ) + + manager.weaviate_manager.add_suggestion.assert_not_called() + + +class TestGenerateStreamAnswer: + @pytest.mark.asyncio + async def test_yields_all_chunks_with_full_text_on_stop(self, manager): + """Streamed tokens are yielded; final stop chunk gets full_text appended.""" + + async def _fake_stream(rag_config, query, context, conversation): + yield {"message": "Hello", "finish_reason": ""} + yield {"message": " world", "finish_reason": "stop"} + + manager.generator_manager.generate_stream = _fake_stream + + rag_config = {"Generator": _make_rag_component_class(selected="OpenAI")} + chunks = [] + async for chunk in manager.generate_stream_answer(rag_config, "q", "ctx", []): + chunks.append(chunk) + + assert len(chunks) == 2 + assert chunks[0]["message"] == "Hello" + assert chunks[1]["message"] == " world" + assert chunks[1]["full_text"] == "Hello world" + + @pytest.mark.asyncio + async def test_yields_nothing_when_stream_is_empty(self, manager): + async def _empty_stream(rag_config, query, context, conversation): + return + yield # make it an async generator + + manager.generator_manager.generate_stream = _empty_stream + + rag_config = {"Generator": _make_rag_component_class(selected="OpenAI")} + chunks = [] + async for chunk in manager.generate_stream_answer(rag_config, "q", "ctx", []): + chunks.append(chunk) + + assert chunks == [] diff --git a/goldenverba/tests/components/weaviate/__init__.py b/goldenverba/tests/components/weaviate/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/goldenverba/tests/components/weaviate/test_weaviate_manager.py b/goldenverba/tests/components/weaviate/test_weaviate_manager.py new file mode 100644 index 0000000000..d4f9d271f9 --- /dev/null +++ b/goldenverba/tests/components/weaviate/test_weaviate_manager.py @@ -0,0 +1,426 @@ +""" +Integration tests for WeaviateManager. + +These tests require a live Weaviate instance. They are skipped automatically +when the required environment variables are not set, so they are safe to run +in any CI environment — they only execute when a real instance is available. + +How to run locally +------------------ +# Against a local Docker instance (docker compose up -d weaviate): + pytest goldenverba/tests/weaviate/ -v + +# Against Weaviate Cloud: + WEAVIATE_TEST_URL=https://my-cluster.weaviate.network \\ + WEAVIATE_TEST_KEY=my-api-key \\ + pytest goldenverba/tests/weaviate/ -v + +Environment variables +--------------------- +WEAVIATE_TEST_URL — cluster URL for Weaviate Cloud. When absent the tests + connect to a local Docker instance at localhost:8080. +WEAVIATE_TEST_KEY — API key. Required when WEAVIATE_TEST_URL is set. + Optional for unauthenticated local instances. +""" + +import os +import pytest +import pytest_asyncio + +# All async tests in this module share a single event loop so that +# module-scoped fixtures (manager, client) remain connected between tests. +pytestmark = pytest.mark.asyncio(loop_scope="module") + +from goldenverba.components.weaviate_manager import WeaviateManager +from goldenverba.components.document import Document +from goldenverba.components.chunk import Chunk + +# --------------------------------------------------------------------------- +# Skip condition +# --------------------------------------------------------------------------- + +# Tests are opt-in: they only run when WEAVIATE_TEST_URL (cloud) or +# WEAVIATE_INTEGRATION=1 (local Docker) is explicitly set. +# This ensures they are safely skipped in CI and local dev by default. +_CLOUD_URL = os.environ.get("WEAVIATE_TEST_URL", "") +_CLOUD_KEY = os.environ.get("WEAVIATE_TEST_KEY", "") +_LOCAL = os.environ.get("WEAVIATE_INTEGRATION", "").lower() in ("1", "true", "yes") + +_ENABLED = bool(_CLOUD_URL or _LOCAL) + +requires_weaviate = pytest.mark.skipif( + not _ENABLED, + reason=( + "Integration tests are opt-in. To run against a local Docker instance: " + "WEAVIATE_INTEGRATION=1 pytest goldenverba/tests/weaviate/ " + "To run against Weaviate Cloud: " + "WEAVIATE_TEST_URL=<url> WEAVIATE_TEST_KEY=<key> pytest goldenverba/tests/weaviate/" + ), +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +# Use a dedicated collection prefix so tests never touch real Verba data. +_TEST_PREFIX = "VERBA_TEST_" + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def manager(): + """A WeaviateManager with test-namespaced collection names.""" + wm = WeaviateManager() + # Redirect all collection names to test-specific names so we never + # interfere with real Verba data in a shared Weaviate instance. + wm.document_collection_name = _TEST_PREFIX + "DOCUMENTS" + wm.config_collection_name = _TEST_PREFIX + "CONFIGURATION" + wm.suggestion_collection_name = _TEST_PREFIX + "SUGGESTIONS" + return wm + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def client(manager): + """ + A live WeaviateAsyncClient. + + Connects to Weaviate Cloud if WEAVIATE_TEST_URL is set, otherwise + falls back to a local Docker instance at localhost:8080. + """ + if _CLOUD_URL: + if not _CLOUD_KEY: + pytest.skip("WEAVIATE_TEST_URL is set but WEAVIATE_TEST_KEY is missing") + c = await manager.connect( + deployment="Weaviate", + weaviateURL=_CLOUD_URL, + weaviateAPIKey=_CLOUD_KEY, + ) + else: + # Local Docker: WEAVIATE_HOST defaults to localhost for tests. + # (In Docker Compose it would be "weaviate", but tests run on the host.) + c = await manager.connect( + deployment="Custom", + weaviateURL=os.environ.get("WEAVIATE_HOST", "localhost"), + weaviateAPIKey="", + port=os.environ.get("WEAVIATE_PORT", "8080"), + ) + yield c + # Teardown: delete every test collection we created so the instance is clean. + for name in list(await c.collections.list_all()): + if name.startswith(_TEST_PREFIX): + await c.collections.delete(name) + await manager.disconnect(c) + + +@pytest_asyncio.fixture(autouse=True, loop_scope="module") +async def clean_collections(manager, client): + """ + Before each test: purge all test collections so tests start from a + known-empty state without depending on execution order. + """ + for name in list(await client.collections.list_all()): + if name.startswith(_TEST_PREFIX): + await client.collections.delete(name) + # Clear the manager's collection verification cache so it re-creates them. + manager._verified_collections.clear() + manager.embedding_table.clear() + manager.cache_table.clear() + yield + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_TEST_EMBEDDER = "test-embedder-model" + + +def make_document(title: str = "Test Doc", n_chunks: int = 3) -> Document: + doc = Document(title=title, content="Hello world", labels=["test"]) + for i in range(n_chunks): + chunk = Chunk( + content=f"chunk {i}", + chunk_id=i, + start_i=i * 10, + end_i=i * 10 + 9, + ) + chunk.vector = [0.1 * (i + 1)] * 4 # tiny fake vector + chunk.pca = [0.0, 0.0, 0.0] + chunk.title = title + chunk.labels = ["test"] + doc.chunks.append(chunk) + return doc + + +# --------------------------------------------------------------------------- +# Connection tests +# --------------------------------------------------------------------------- + + +@requires_weaviate +class TestConnection: + async def test_client_is_ready(self, client): + assert await client.is_ready() + + async def test_invalid_deployment_raises(self, manager): + with pytest.raises(Exception, match="Invalid deployment type"): + await manager.connect( + deployment="NonExistent", + weaviateURL="", + weaviateAPIKey="", + ) + + async def test_cloud_missing_key_raises(self, manager): + with pytest.raises(Exception): + await manager.connect( + deployment="Weaviate", + weaviateURL="https://example.weaviate.network", + weaviateAPIKey="", # empty — should raise + ) + + +# --------------------------------------------------------------------------- +# Collection management +# --------------------------------------------------------------------------- + + +@requires_weaviate +class TestCollections: + async def test_verify_collection_creates_if_absent(self, manager, client): + name = _TEST_PREFIX + "NEWCOL" + assert not await client.collections.exists(name) + result = await manager.verify_collection(client, name) + assert result is True + assert await client.collections.exists(name) + + async def test_verify_collection_caches_result(self, manager, client): + name = _TEST_PREFIX + "CACHED" + await manager.verify_collection(client, name) + client_id = id(client) + assert name in manager._verified_collections.get(client_id, set()) + + # Delete the collection externally — the cache still says "exists" + await client.collections.delete(name) + # Second call should return True from cache without hitting Weaviate + result = await manager.verify_collection(client, name) + assert result is True + + async def test_verify_embedding_collection(self, manager, client): + result = await manager.verify_embedding_collection(client, _TEST_EMBEDDER) + assert result is True + assert _TEST_EMBEDDER in manager.embedding_table + collection_name = manager.embedding_table[_TEST_EMBEDDER] + assert await client.collections.exists(collection_name) + + async def test_verify_embedding_collection_cached_on_second_call(self, manager, client): + await manager.verify_embedding_collection(client, _TEST_EMBEDDER) + # Second call: embedder is already in embedding_table, returns True immediately + result = await manager.verify_embedding_collection(client, _TEST_EMBEDDER) + assert result is True + + +# --------------------------------------------------------------------------- +# Configuration CRUD +# --------------------------------------------------------------------------- + + +@requires_weaviate +class TestConfig: + _UUID = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + + async def test_set_and_get_config(self, manager, client): + config = {"key": "value", "nested": {"a": 1}} + await manager.set_config(client, self._UUID, config) + loaded = await manager.get_config(client, self._UUID) + assert loaded == config + + async def test_get_config_missing_returns_none(self, manager, client): + result = await manager.get_config(client, self._UUID) + assert result is None + + async def test_set_config_overwrites(self, manager, client): + await manager.set_config(client, self._UUID, {"v": 1}) + await manager.set_config(client, self._UUID, {"v": 2}) + loaded = await manager.get_config(client, self._UUID) + assert loaded == {"v": 2} + + async def test_reset_config(self, manager, client): + await manager.set_config(client, self._UUID, {"v": 1}) + await manager.reset_config(client, self._UUID) + result = await manager.get_config(client, self._UUID) + assert result is None + + +# --------------------------------------------------------------------------- +# Document CRUD +# --------------------------------------------------------------------------- + + +@requires_weaviate +class TestDocuments: + async def test_exist_document_name_false_when_empty(self, manager, client): + result = await manager.exist_document_name(client, "Nonexistent") + assert result is None + + async def test_import_and_exist(self, manager, client): + doc = make_document("Alpha") + doc.meta = {"Embedder": {"config": {"Model": {"value": _TEST_EMBEDDER}}}} + await manager.import_document(client, doc, _TEST_EMBEDDER) + + uuid = await manager.exist_document_name(client, "Alpha") + assert uuid is not None + + async def test_get_document(self, manager, client): + doc = make_document("Beta") + doc.meta = {"Embedder": {"config": {"Model": {"value": _TEST_EMBEDDER}}}} + await manager.import_document(client, doc, _TEST_EMBEDDER) + + uuid = await manager.exist_document_name(client, "Beta") + fetched = await manager.get_document(client, str(uuid)) + assert fetched is not None + assert fetched["title"] == "Beta" + + async def test_get_document_missing_returns_none(self, manager, client): + result = await manager.get_document(client, "00000000-0000-0000-0000-000000000000") + assert result is None + + async def test_delete_document(self, manager, client): + doc = make_document("Gamma") + doc.meta = {"Embedder": {"config": {"Model": {"value": _TEST_EMBEDDER}}}} + await manager.import_document(client, doc, _TEST_EMBEDDER) + + uuid = await manager.exist_document_name(client, "Gamma") + assert uuid is not None + + await manager.delete_document(client, str(uuid)) + assert await manager.exist_document_name(client, "Gamma") is None + + async def test_delete_all_documents(self, manager, client): + for title in ["D1", "D2", "D3"]: + doc = make_document(title) + doc.meta = {"Embedder": {"config": {"Model": {"value": _TEST_EMBEDDER}}}} + await manager.import_document(client, doc, _TEST_EMBEDDER) + + await manager.delete_all_documents(client) + + for title in ["D1", "D2", "D3"]: + assert await manager.exist_document_name(client, title) is None + + async def test_get_documents_pagination(self, manager, client): + for i in range(5): + doc = make_document(f"Page{i}") + doc.meta = {"Embedder": {"config": {"Model": {"value": _TEST_EMBEDDER}}}} + await manager.import_document(client, doc, _TEST_EMBEDDER) + + results, total = await manager.get_documents( + client, query="", pageSize=3, page=1, labels=[] + ) + assert total == 5 + assert len(results) == 3 + + async def test_get_documents_bm25_search(self, manager, client): + for title in ["Python guide", "JavaScript guide", "Rust primer"]: + doc = make_document(title) + doc.meta = {"Embedder": {"config": {"Model": {"value": _TEST_EMBEDDER}}}} + await manager.import_document(client, doc, _TEST_EMBEDDER) + + results, _ = await manager.get_documents( + client, query="Python", pageSize=10, page=1, labels=[] + ) + titles = [r["title"] for r in results] + assert "Python guide" in titles + + async def test_get_labels(self, manager, client): + doc = make_document("Labelled") + doc.labels = ["science", "tech"] + doc.meta = {"Embedder": {"config": {"Model": {"value": _TEST_EMBEDDER}}}} + await manager.import_document(client, doc, _TEST_EMBEDDER) + + labels = await manager.get_labels(client) + assert "science" in labels + assert "tech" in labels + + +# --------------------------------------------------------------------------- +# Chunks +# --------------------------------------------------------------------------- + + +@requires_weaviate +class TestChunks: + async def test_get_chunk(self, manager, client): + doc = make_document("ChunkDoc", n_chunks=3) + doc.meta = {"Embedder": {"config": {"Model": {"value": _TEST_EMBEDDER}}}} + await manager.import_document(client, doc, _TEST_EMBEDDER) + + uuid = await manager.exist_document_name(client, "ChunkDoc") + fetched_doc = await manager.get_document(client, str(uuid)) + + # get_chunk_by_ids: fetch chunks 0 and 1 + chunks = await manager.get_chunk_by_ids(client, _TEST_EMBEDDER, str(uuid), [0, 1]) + assert len(chunks) == 2 + chunk_ids = {c.properties["chunk_id"] for c in chunks} + assert chunk_ids == {0, 1} + + async def test_get_chunk_count(self, manager, client): + doc = make_document("CountDoc", n_chunks=4) + doc.meta = {"Embedder": {"config": {"Model": {"value": _TEST_EMBEDDER}}}} + await manager.import_document(client, doc, _TEST_EMBEDDER) + + uuid = await manager.exist_document_name(client, "CountDoc") + count = await manager.get_chunk_count(client, _TEST_EMBEDDER, str(uuid)) + assert count == 4 + + async def test_hybrid_chunks_returns_results(self, manager, client): + doc = make_document("HybridDoc", n_chunks=2) + doc.meta = {"Embedder": {"config": {"Model": {"value": _TEST_EMBEDDER}}}} + await manager.import_document(client, doc, _TEST_EMBEDDER) + + # Hybrid search with a tiny random vector — should return chunks + results = await manager.hybrid_chunks( + client, + embedder=_TEST_EMBEDDER, + query="chunk", + vector=[0.1, 0.1, 0.1, 0.1], + limit_mode="Limit", + limit=10, + labels=[], + document_uuids=[], + ) + assert len(results) > 0 + + +# --------------------------------------------------------------------------- +# Suggestions +# --------------------------------------------------------------------------- + + +@requires_weaviate +class TestSuggestions: + async def test_add_and_retrieve_suggestion(self, manager, client): + await manager.add_suggestion(client, "what is RAG?") + results = await manager.retrieve_suggestions(client, "RAG", limit=5) + assert any(r["query"] == "what is RAG?" for r in results) + + async def test_add_suggestion_deduplicates(self, manager, client): + await manager.add_suggestion(client, "duplicate query") + await manager.add_suggestion(client, "duplicate query") + results = await manager.retrieve_suggestions(client, "duplicate", limit=10) + matches = [r for r in results if r["query"] == "duplicate query"] + assert len(matches) == 1 + + async def test_delete_suggestion(self, manager, client): + await manager.add_suggestion(client, "to be deleted") + results = await manager.retrieve_suggestions(client, "deleted", limit=5) + uuid = next(r["uuid"] for r in results if r["query"] == "to be deleted") + + await manager.delete_suggestions(client, uuid) + results_after = await manager.retrieve_suggestions(client, "deleted", limit=5) + assert not any(r["query"] == "to be deleted" for r in results_after) + + async def test_retrieve_all_suggestions_pagination(self, manager, client): + for i in range(5): + await manager.add_suggestion(client, f"unique suggestion {i}") + + page1, total = await manager.retrieve_all_suggestions(client, page=1, pageSize=3) + assert total >= 5 + assert len(page1) == 3 diff --git a/goldenverba/tests/server/__init__.py b/goldenverba/tests/server/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/goldenverba/tests/server/test_api.py b/goldenverba/tests/server/test_api.py new file mode 100644 index 0000000000..10126d0676 --- /dev/null +++ b/goldenverba/tests/server/test_api.py @@ -0,0 +1,532 @@ +""" +Tests for goldenverba/server/api.py FastAPI endpoints. + +Strategy: patch the module-level `manager` and `client_manager` singletons +so that no real Weaviate connection or LLM call is required. +""" + +import pytest +from fastapi.testclient import TestClient +from unittest.mock import AsyncMock, MagicMock, patch + +import goldenverba.server.api as api_module +from weaviate.client import WeaviateAsyncClient + +# --------------------------------------------------------------------------- +# Shared payload helpers +# --------------------------------------------------------------------------- + +CREDENTIALS = {"deployment": "Local", "url": "", "key": ""} + + +def creds_payload(**extra): + return {**CREDENTIALS, **extra} + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_weaviate_client(): + # spec=WeaviateAsyncClient makes isinstance() checks in api.py pass + return MagicMock(spec=WeaviateAsyncClient) + + +@pytest.fixture +def mock_manager(): + m = MagicMock() + m.get_deployments = AsyncMock( + return_value={"WEAVIATE_URL_VERBA": "", "WEAVIATE_API_KEY_VERBA": ""} + ) + m.load_rag_config = AsyncMock( + return_value={ + "Reader": {"components": {}, "selected": ""}, + "Chunker": {"components": {}, "selected": ""}, + "Embedder": {"components": {}, "selected": ""}, + "Retriever": {"components": {}, "selected": ""}, + "Generator": {"components": {}, "selected": ""}, + } + ) + m.load_user_config = AsyncMock(return_value={"getting_started": False}) + m.load_theme_config = AsyncMock(return_value=(None, None)) + m.set_rag_config = AsyncMock() + m.set_user_config = AsyncMock() + m.set_theme_config = AsyncMock() + m.retrieve_chunks = AsyncMock(return_value=([], "")) + m.get_content = AsyncMock(return_value=([], 0)) + + wm = MagicMock() + wm.get_document = AsyncMock( + return_value={ + "title": "test.txt", + "content": "", + "extension": ".txt", + "fileSize": 42, + "labels": [], + "source": "", + "meta": "{}", + "metadata": "", + } + ) + wm.get_documents = AsyncMock(return_value=([], 0)) + wm.get_labels = AsyncMock(return_value=[]) + wm.get_chunks = AsyncMock(return_value=[]) + wm.get_chunk = AsyncMock(return_value={}) + wm.get_vectors = AsyncMock( + return_value={"embedder": "test", "dimensions": 3, "groups": []} + ) + wm.get_datacount = AsyncMock(return_value=0) + wm.get_metadata = AsyncMock( + return_value=( + {"node_count": 1, "weaviate_version": "1.0", "nodes": []}, + {"collection_count": 0, "collections": []}, + ) + ) + wm.delete_document = AsyncMock() + wm.delete_all = AsyncMock() + wm.delete_all_documents = AsyncMock() + wm.delete_all_configs = AsyncMock() + wm.delete_all_suggestions = AsyncMock() + wm.retrieve_suggestions = AsyncMock(return_value=[]) + wm.retrieve_all_suggestions = AsyncMock(return_value=([], 0)) + wm.delete_suggestions = AsyncMock() + m.weaviate_manager = wm + return m + + +@pytest.fixture +def mock_client_manager(mock_weaviate_client): + m = MagicMock() + m.connect = AsyncMock(return_value=mock_weaviate_client) + m.disconnect = AsyncMock() + m.clean_up = AsyncMock() + return m + + +@pytest.fixture +def client(mock_manager, mock_client_manager): + # The same-origin middleware checks that Origin matches the server's base URL. + # TestClient uses http://testserver, so we pass that as the default Origin header. + with ( + patch.object(api_module, "manager", mock_manager), + patch.object(api_module, "client_manager", mock_client_manager), + ): + with TestClient( + api_module.app, headers={"origin": "http://testserver"} + ) as c: + yield c + + +# --------------------------------------------------------------------------- +# /api/health +# --------------------------------------------------------------------------- + + +class TestHealthEndpoint: + def test_health_returns_200(self, client): + resp = client.get("/api/health") + assert resp.status_code == 200 + + def test_health_body_has_message(self, client): + data = client.get("/api/health").json() + assert data["message"] == "Alive!" + + def test_health_calls_clean_up(self, client, mock_client_manager): + client.get("/api/health") + mock_client_manager.clean_up.assert_awaited_once() + + def test_health_contains_production_field(self, client): + data = client.get("/api/health").json() + assert "production" in data + + def test_health_contains_deployments(self, client): + data = client.get("/api/health").json() + assert "deployments" in data + + +# --------------------------------------------------------------------------- +# /api/connect +# --------------------------------------------------------------------------- + + +class TestConnectEndpoint: + def _payload(self, **extra): + return {"credentials": CREDENTIALS, "port": "8080", **extra} + + def test_connect_success_returns_200(self, client): + resp = client.post("/api/connect", json=self._payload()) + assert resp.status_code == 200 + + def test_connect_success_connected_true(self, client): + data = client.post("/api/connect", json=self._payload()).json() + assert data["connected"] is True + + def test_connect_success_returns_rag_config(self, client): + data = client.post("/api/connect", json=self._payload()).json() + assert "rag_config" in data + + def test_connect_failure_returns_400(self, client, mock_client_manager): + mock_client_manager.connect = AsyncMock( + side_effect=Exception("connection refused") + ) + resp = client.post("/api/connect", json=self._payload()) + assert resp.status_code == 400 + + def test_connect_failure_connected_false(self, client, mock_client_manager): + mock_client_manager.connect = AsyncMock( + side_effect=Exception("connection refused") + ) + data = client.post("/api/connect", json=self._payload()).json() + assert data["connected"] is False + + def test_connect_failure_contains_error(self, client, mock_client_manager): + mock_client_manager.connect = AsyncMock( + side_effect=Exception("bad credentials") + ) + data = client.post("/api/connect", json=self._payload()).json() + assert "error" in data + assert data["error"] != "" + + +# --------------------------------------------------------------------------- +# /api/get_rag_config +# --------------------------------------------------------------------------- + + +class TestGetRagConfig: + def test_returns_200(self, client): + resp = client.post("/api/get_rag_config", json=CREDENTIALS) + assert resp.status_code == 200 + + def test_returns_rag_config(self, client): + data = client.post("/api/get_rag_config", json=CREDENTIALS).json() + assert "rag_config" in data + + def test_error_returns_500(self, client, mock_manager): + mock_manager.load_rag_config = AsyncMock(side_effect=Exception("db error")) + resp = client.post("/api/get_rag_config", json=CREDENTIALS) + assert resp.status_code == 500 + + +# --------------------------------------------------------------------------- +# /api/set_rag_config +# --------------------------------------------------------------------------- + + +class TestSetRagConfig: + def _payload(self): + rag_config = { + "Reader": {"selected": "", "components": {}}, + "Chunker": {"selected": "", "components": {}}, + "Embedder": {"selected": "", "components": {}}, + "Retriever": {"selected": "", "components": {}}, + "Generator": {"selected": "", "components": {}}, + } + return {"rag_config": rag_config, "credentials": CREDENTIALS} + + def test_success_returns_200_status(self, client): + data = client.post("/api/set_rag_config", json=self._payload()).json() + assert data["status"] == 200 + + def test_calls_set_rag_config(self, client, mock_manager): + client.post("/api/set_rag_config", json=self._payload()) + mock_manager.set_rag_config.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# /api/query +# --------------------------------------------------------------------------- + + +class TestQueryEndpoint: + def _payload(self, query="what is Verba?"): + return { + "query": query, + "RAG": { + "Reader": {"selected": "", "components": {}}, + "Chunker": {"selected": "", "components": {}}, + "Embedder": {"selected": "", "components": {}}, + "Retriever": {"selected": "", "components": {}}, + "Generator": {"selected": "", "components": {}}, + }, + "labels": [], + "documentFilter": [], + "credentials": CREDENTIALS, + } + + def test_returns_200(self, client): + resp = client.post("/api/query", json=self._payload()) + assert resp.status_code == 200 + + def test_returns_documents_and_context(self, client): + data = client.post("/api/query", json=self._payload()).json() + assert "documents" in data + assert "context" in data + + def test_error_path_returns_error_field(self, client, mock_manager): + mock_manager.retrieve_chunks = AsyncMock(side_effect=Exception("retrieval failed")) + data = client.post("/api/query", json=self._payload()).json() + assert data["error"] != "" + assert data["documents"] == [] + + def test_query_too_long_rejected(self, client): + resp = client.post("/api/query", json=self._payload(query="x" * 50_001)) + assert resp.status_code == 422 + + +# --------------------------------------------------------------------------- +# /api/get_document +# --------------------------------------------------------------------------- + + +class TestGetDocument: + def _payload(self, uuid="abc-123"): + return {"uuid": uuid, "credentials": CREDENTIALS} + + def test_returns_document(self, client): + data = client.post("/api/get_document", json=self._payload()).json() + assert data["document"]["title"] == "test.txt" + + def test_document_not_found_returns_error(self, client, mock_manager): + mock_manager.weaviate_manager.get_document = AsyncMock(return_value=None) + data = client.post("/api/get_document", json=self._payload()).json() + assert data["document"] is None + assert data["error"] != "" + + def test_exception_returns_error(self, client, mock_manager): + mock_manager.weaviate_manager.get_document = AsyncMock( + side_effect=Exception("weaviate down") + ) + data = client.post("/api/get_document", json=self._payload()).json() + assert data["document"] is None + + +# --------------------------------------------------------------------------- +# /api/get_all_documents +# --------------------------------------------------------------------------- + + +class TestGetAllDocuments: + def _payload(self, query="", page=1, page_size=10): + return { + "query": query, + "labels": [], + "page": page, + "pageSize": page_size, + "credentials": CREDENTIALS, + } + + def test_returns_200(self, client): + resp = client.post("/api/get_all_documents", json=self._payload()) + assert resp.status_code == 200 + + def test_returns_empty_list_when_no_docs(self, client): + data = client.post("/api/get_all_documents", json=self._payload()).json() + assert data["documents"] == [] + assert data["totalDocuments"] == 0 + + def test_returns_documents_from_manager(self, client, mock_manager): + mock_manager.weaviate_manager.get_documents = AsyncMock( + return_value=([{"title": "foo.txt", "uuid": "u1", "labels": []}], 1) + ) + data = client.post("/api/get_all_documents", json=self._payload()).json() + assert len(data["documents"]) == 1 + assert data["totalDocuments"] == 1 + + +# --------------------------------------------------------------------------- +# /api/delete_document +# --------------------------------------------------------------------------- + + +class TestDeleteDocument: + def _payload(self, uuid="abc-123"): + return {"uuid": uuid, "credentials": CREDENTIALS} + + def test_returns_200_on_success(self, client): + resp = client.post("/api/delete_document", json=self._payload()) + assert resp.status_code == 200 + + def test_calls_delete_document(self, client, mock_manager): + client.post("/api/delete_document", json=self._payload("my-uuid")) + mock_manager.weaviate_manager.delete_document.assert_awaited_once() + + def test_exception_returns_400(self, client, mock_manager): + mock_manager.weaviate_manager.delete_document = AsyncMock( + side_effect=Exception("weaviate down") + ) + resp = client.post("/api/delete_document", json=self._payload()) + assert resp.status_code == 400 + + +# --------------------------------------------------------------------------- +# /api/reset +# --------------------------------------------------------------------------- + + +class TestResetEndpoint: + def _payload(self, mode="ALL"): + return {"resetMode": mode, "credentials": CREDENTIALS} + + def test_reset_all_calls_delete_all(self, client, mock_manager): + client.post("/api/reset", json=self._payload("ALL")) + mock_manager.weaviate_manager.delete_all.assert_awaited_once() + + def test_reset_documents_calls_delete_all_documents(self, client, mock_manager): + client.post("/api/reset", json=self._payload("DOCUMENTS")) + mock_manager.weaviate_manager.delete_all_documents.assert_awaited_once() + + def test_reset_config_calls_delete_all_configs(self, client, mock_manager): + client.post("/api/reset", json=self._payload("CONFIG")) + mock_manager.weaviate_manager.delete_all_configs.assert_awaited_once() + + def test_reset_suggestions_calls_delete_all_suggestions(self, client, mock_manager): + client.post("/api/reset", json=self._payload("SUGGESTIONS")) + mock_manager.weaviate_manager.delete_all_suggestions.assert_awaited_once() + + def test_exception_returns_500(self, client, mock_manager): + mock_manager.weaviate_manager.delete_all = AsyncMock( + side_effect=Exception("weaviate down") + ) + resp = client.post("/api/reset", json=self._payload("ALL")) + assert resp.status_code == 500 + + +# --------------------------------------------------------------------------- +# /api/get_meta +# --------------------------------------------------------------------------- + + +class TestGetMeta: + def test_returns_node_and_collection_payload(self, client): + data = client.post("/api/get_meta", json=CREDENTIALS).json() + assert "node_payload" in data + assert "collection_payload" in data + + def test_exception_returns_error(self, client, mock_manager): + mock_manager.weaviate_manager.get_metadata = AsyncMock( + side_effect=Exception("metadata unavailable") + ) + data = client.post("/api/get_meta", json=CREDENTIALS).json() + assert data["error"] != "" + + +# --------------------------------------------------------------------------- +# /api/get_suggestions /api/delete_suggestion +# --------------------------------------------------------------------------- + + +class TestSuggestions: + def test_get_suggestions_returns_list(self, client): + data = client.post( + "/api/get_suggestions", + json={"query": "test", "limit": 5, "credentials": CREDENTIALS}, + ).json() + assert "suggestions" in data + assert isinstance(data["suggestions"], list) + + def test_get_suggestions_exception_returns_empty(self, client, mock_manager): + mock_manager.weaviate_manager.retrieve_suggestions = AsyncMock( + side_effect=Exception("boom") + ) + data = client.post( + "/api/get_suggestions", + json={"query": "test", "limit": 5, "credentials": CREDENTIALS}, + ).json() + assert data["suggestions"] == [] + + def test_delete_suggestion_returns_200_status(self, client): + data = client.post( + "/api/delete_suggestion", + json={"uuid": "some-uuid", "credentials": CREDENTIALS}, + ).json() + assert data["status"] == 200 + + def test_get_all_suggestions_returns_list_and_count(self, client): + data = client.post( + "/api/get_all_suggestions", + json={"page": 1, "pageSize": 10, "credentials": CREDENTIALS}, + ).json() + assert "suggestions" in data + assert "total_count" in data + + +# --------------------------------------------------------------------------- +# /api/get_content +# --------------------------------------------------------------------------- + + +class TestGetContent: + def _payload(self): + return { + "uuid": "doc-uuid", + "page": 1, + "chunkScores": [], + "credentials": CREDENTIALS, + } + + def test_returns_content_and_max_page(self, client): + data = client.post("/api/get_content", json=self._payload()).json() + assert "content" in data + assert "maxPage" in data + + def test_exception_returns_error(self, client, mock_manager): + mock_manager.get_content = AsyncMock(side_effect=Exception("db error")) + data = client.post("/api/get_content", json=self._payload()).json() + assert data["error"] != "" + + +# --------------------------------------------------------------------------- +# /api/get_datacount +# --------------------------------------------------------------------------- + + +class TestGetDatacount: + def test_returns_datacount(self, client): + data = client.post( + "/api/get_datacount", + json={ + "embedding_model": "all-MiniLM-L6-v2", + "documentFilter": [], + "credentials": CREDENTIALS, + }, + ).json() + assert "datacount" in data + assert data["datacount"] == 0 + + def test_exception_returns_zero(self, client, mock_manager): + mock_manager.weaviate_manager.get_datacount = AsyncMock( + side_effect=Exception("boom") + ) + data = client.post( + "/api/get_datacount", + json={ + "embedding_model": "all-MiniLM-L6-v2", + "documentFilter": [], + "credentials": CREDENTIALS, + }, + ).json() + assert data["datacount"] == 0 + + +# --------------------------------------------------------------------------- +# /api/get_labels +# --------------------------------------------------------------------------- + + +class TestGetLabels: + def test_returns_labels_from_manager(self, client, mock_manager): + mock_manager.weaviate_manager.get_labels = AsyncMock( + return_value=["finance", "legal"] + ) + data = client.post("/api/get_labels", json=CREDENTIALS).json() + assert data["labels"] == ["finance", "legal"] + + def test_exception_returns_empty_list(self, client, mock_manager): + mock_manager.weaviate_manager.get_labels = AsyncMock( + side_effect=Exception("boom") + ) + data = client.post("/api/get_labels", json=CREDENTIALS).json() + assert data["labels"] == [] diff --git a/goldenverba/tests/server/test_helpers.py b/goldenverba/tests/server/test_helpers.py new file mode 100644 index 0000000000..b6d7b0e58c --- /dev/null +++ b/goldenverba/tests/server/test_helpers.py @@ -0,0 +1,450 @@ +""" +Tests for goldenverba/server/helpers.py + +Covers BatchManager and LoggerManager. +No real network calls or Weaviate connections are made. +""" +import json +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from goldenverba.server.helpers import BatchManager, LoggerManager +from goldenverba.server.types import ( + DataBatchPayload, + FileStatus, + Credentials, + FileConfig, + RAGComponentClass, +) + + +# --------------------------------------------------------------------------- +# Helpers / fixtures +# --------------------------------------------------------------------------- + +DUMMY_CREDENTIALS = Credentials( + deployment="Local", + url="http://localhost:8080", + key="", +) + + +def _make_file_config( + file_id: str = "file-001", + filename: str = "test.txt", +) -> FileConfig: + """Return a minimal but valid FileConfig object.""" + return FileConfig( + fileID=file_id, + filename=filename, + isURL=False, + overwrite=False, + extension=".txt", + source="local", + content="Hello world", + labels=[], + rag_config={}, + file_size=11, + status=FileStatus.READY, + metadata="", + status_report={}, + ) + + +def _serialise_file_config(fc: FileConfig) -> str: + """Return the JSON string that the frontend would send as a batch payload.""" + return fc.model_dump_json() + + +def _make_payload( + file_id: str, + chunk: str, + order: int, + total: int, + is_last: bool = False, +) -> DataBatchPayload: + return DataBatchPayload( + fileID=file_id, + chunk=chunk, + order=order, + total=total, + isLastChunk=is_last, + credentials=DUMMY_CREDENTIALS, + ) + + +def _split_into_chunks(text: str, n: int) -> list[str]: + """Split *text* into *n* roughly-equal string chunks.""" + size = max(1, len(text) // n) + parts = [text[i : i + size] for i in range(0, len(text), size)] + # If rounding produced more than n parts, merge the tail into the last part + while len(parts) > n: + parts[-2] = parts[-2] + parts[-1] + parts.pop() + return parts + + +# --------------------------------------------------------------------------- +# BatchManager tests +# --------------------------------------------------------------------------- + + +class TestBatchManagerNormalFlow: + """Single-chunk and multi-chunk happy-path scenarios.""" + + def test_single_chunk_returns_file_config(self): + """A batch with total=1 should resolve immediately.""" + manager = BatchManager() + fc = _make_file_config() + payload = _make_payload( + file_id=fc.fileID, + chunk=_serialise_file_config(fc), + order=0, + total=1, + is_last=True, + ) + + result = manager.add_batch(payload) + + assert result is not None, "Expected a FileConfig back from a complete batch" + assert isinstance(result, FileConfig) + assert result.fileID == fc.fileID + assert result.filename == fc.filename + + def test_batch_removed_after_assembly(self): + """Once all chunks arrive the entry must be cleaned up.""" + manager = BatchManager() + fc = _make_file_config() + full_json = _serialise_file_config(fc) + parts = _split_into_chunks(full_json, 3) + + for i, part in enumerate(parts): + is_last = i == len(parts) - 1 + manager.add_batch( + _make_payload(fc.fileID, part, i, len(parts), is_last) + ) + + assert fc.fileID not in manager.batches, ( + "Batch entry must be removed once the FileConfig is assembled" + ) + + def test_multi_chunk_reassembly(self): + """Splitting a FileConfig JSON into multiple ordered chunks then reassembling + must produce the same FileConfig.""" + manager = BatchManager() + fc = _make_file_config(file_id="reassembly-test", filename="reassembly.txt") + full_json = _serialise_file_config(fc) + parts = _split_into_chunks(full_json, 4) + + result = None + for i, part in enumerate(parts): + is_last = i == len(parts) - 1 + result = manager.add_batch( + _make_payload(fc.fileID, part, i, len(parts), is_last) + ) + + assert result is not None + assert result.fileID == fc.fileID + assert result.filename == fc.filename + assert result.extension == fc.extension + + def test_incomplete_batch_returns_none(self): + """Sending fewer chunks than `total` must not resolve yet.""" + manager = BatchManager() + fc = _make_file_config() + full_json = _serialise_file_config(fc) + parts = _split_into_chunks(full_json, 3) + + # Only send the first chunk, not the remaining two + result = manager.add_batch( + _make_payload(fc.fileID, parts[0], 0, 3, is_last=False) + ) + + assert result is None, "Batch should not resolve before all chunks arrive" + + def test_batch_entry_exists_while_incomplete(self): + """An in-progress batch must remain in the dict.""" + manager = BatchManager() + fc = _make_file_config() + full_json = _serialise_file_config(fc) + parts = _split_into_chunks(full_json, 3) + + manager.add_batch(_make_payload(fc.fileID, parts[0], 0, 3, is_last=False)) + + assert fc.fileID in manager.batches, ( + "In-progress batch must remain registered in BatchManager" + ) + + +class TestBatchManagerOutOfOrderChunks: + """Chunks that arrive in a non-sequential order. + + NOTE: BatchManager.check_batch() joins chunks using dict insertion order, + not sorted key order. Sending chunks out of sequence therefore produces a + garbled JSON string and the assembly fails silently (returns None / raises). + These tests document the *actual* behaviour so that a future fix can be + verified by updating the assertions here. + """ + + def test_in_order_chunks_assemble_correctly(self): + """Chunks delivered in order 0→1→2 must produce a valid FileConfig.""" + manager = BatchManager() + fc = _make_file_config(file_id="in-order-test") + full_json = _serialise_file_config(fc) + parts = _split_into_chunks(full_json, 3) + + result = None + for i, part in enumerate(parts): + is_last = i == len(parts) - 1 + result = manager.add_batch( + _make_payload(fc.fileID, part, i, len(parts), is_last) + ) + + assert result is not None, "In-order chunks must assemble correctly" + assert result.fileID == fc.fileID + + def test_out_of_order_single_chunk_resolves(self): + """A single-chunk batch is always in-order; it must resolve regardless.""" + manager = BatchManager() + fc = _make_file_config(file_id="single-chunk-ooo") + full_json = _serialise_file_config(fc) + + result = manager.add_batch( + _make_payload(fc.fileID, full_json, 0, 1, is_last=True) + ) + + assert result is not None + assert result.fileID == fc.fileID + + def test_two_chunks_in_order_assemble(self): + """Two chunks delivered in the correct order must produce a FileConfig.""" + manager = BatchManager() + fc = _make_file_config(file_id="two-chunk-test") + full_json = _serialise_file_config(fc) + parts = _split_into_chunks(full_json, 2) + + manager.add_batch(_make_payload(fc.fileID, parts[0], 0, 2, is_last=False)) + result = manager.add_batch( + _make_payload(fc.fileID, parts[1], 1, 2, is_last=True) + ) + + assert result is not None, "Two in-order chunks should assemble correctly" + assert result.fileID == fc.fileID + + +class TestBatchManagerTTL: + """Stale entries are evicted once their TTL expires.""" + + def test_stale_entry_evicted_on_next_add(self): + """An entry older than _BATCH_TTL_SECONDS is removed the next time add_batch + is called for a *different* file.""" + import time + from goldenverba.server.helpers import _BATCH_TTL_SECONDS + + manager = BatchManager() + old_fc = _make_file_config(file_id="old-upload") + new_fc = _make_file_config(file_id="new-upload") + + # Seed the old entry + manager.add_batch(_make_payload(old_fc.fileID, "x", 0, 2, is_last=False)) + + # Back-date its created_at so it looks expired + manager.batches[old_fc.fileID]["created_at"] = ( + time.monotonic() - _BATCH_TTL_SECONDS - 1 + ) + + # Trigger eviction via a new batch + manager.add_batch( + _make_payload(new_fc.fileID, _serialise_file_config(new_fc), 0, 1, is_last=True) + ) + + assert old_fc.fileID not in manager.batches, ( + "Stale entry must be evicted after TTL" + ) + + def test_fresh_entry_not_evicted(self): + """An entry within its TTL must not be evicted.""" + manager = BatchManager() + fc = _make_file_config(file_id="fresh-upload") + other_fc = _make_file_config(file_id="trigger-upload") + + manager.add_batch(_make_payload(fc.fileID, "x", 0, 2, is_last=False)) + # Trigger eviction check without ageing the entry + manager.add_batch( + _make_payload(other_fc.fileID, _serialise_file_config(other_fc), 0, 1, is_last=True) + ) + + assert fc.fileID in manager.batches, "Fresh entry must survive eviction check" + + +class TestBatchManagerAbandonedBatch: + """A partial upload that never completes stays in the dict until TTL.""" + + def test_abandoned_batch_stays_in_dict(self): + """A partial batch whose remaining chunks never arrive must not disappear.""" + manager = BatchManager() + fc = _make_file_config(file_id="abandoned") + full_json = _serialise_file_config(fc) + parts = _split_into_chunks(full_json, 5) + + # Deliver only 2 of 5 chunks, then stop + manager.add_batch(_make_payload(fc.fileID, parts[0], 0, 5, is_last=False)) + manager.add_batch(_make_payload(fc.fileID, parts[1], 1, 5, is_last=False)) + + # Entry must still be present — TTL hasn't elapsed (300 s) + assert fc.fileID in manager.batches + assert len(manager.batches[fc.fileID]["chunks"]) == 2 + + def test_multiple_independent_batches_tracked_separately(self): + """Two concurrent uploads must not interfere with each other.""" + manager = BatchManager() + fc1 = _make_file_config(file_id="file-A", filename="a.txt") + fc2 = _make_file_config(file_id="file-B", filename="b.txt") + + json1 = _serialise_file_config(fc1) + json2 = _serialise_file_config(fc2) + + # Start both uploads, complete only the second + manager.add_batch(_make_payload(fc1.fileID, json1[:10], 0, 2, is_last=False)) + manager.add_batch( + _make_payload(fc2.fileID, json2, 0, 1, is_last=True) + ) + + # fc1 should still be pending; fc2 should be gone (completed) + assert fc1.fileID in manager.batches + assert fc2.fileID not in manager.batches + + +class TestBatchManagerDuplicateChunk: + """Sending the same order index twice must overwrite the stored chunk.""" + + def test_duplicate_order_overwrites(self): + """The second delivery of order=0 replaces the first.""" + manager = BatchManager() + fc = _make_file_config() + full_json = _serialise_file_config(fc) + + # Send a garbage chunk at order=0 first, then the correct one + manager.add_batch(_make_payload(fc.fileID, "GARBAGE", 0, 1, is_last=False)) + result = manager.add_batch( + _make_payload(fc.fileID, full_json, 0, 1, is_last=True) + ) + + assert result is not None + assert result.fileID == fc.fileID + + +class TestBatchManagerLastChunkFlag: + """isLastChunk=True on a partial batch should clean up without assembling.""" + + def test_is_last_chunk_triggers_cleanup_even_if_incomplete(self): + """When isLastChunk is True the entry is removed regardless of completeness.""" + manager = BatchManager() + fc = _make_file_config(file_id="last-flag-test") + full_json = _serialise_file_config(fc) + parts = _split_into_chunks(full_json, 3) + + # Send only one chunk but mark it as the last + manager.add_batch( + _make_payload(fc.fileID, parts[0], 0, 3, is_last=True) + ) + + assert fc.fileID not in manager.batches, ( + "isLastChunk=True must remove the entry even if not all chunks arrived" + ) + + +# --------------------------------------------------------------------------- +# LoggerManager tests +# --------------------------------------------------------------------------- + + +class TestLoggerManagerInstantiation: + def test_instantiation_without_socket(self): + """LoggerManager must be constructable without a socket.""" + lm = LoggerManager() + assert lm.socket is None + + def test_instantiation_with_socket(self): + """LoggerManager must store the provided socket.""" + mock_socket = MagicMock() + lm = LoggerManager(socket=mock_socket) + assert lm.socket is mock_socket + + +class TestLoggerManagerSendReport: + @pytest.mark.asyncio + async def test_send_report_without_socket_does_not_crash(self): + """send_report must not raise even when no socket is configured.""" + lm = LoggerManager() + # Should complete without raising + await lm.send_report( + file_Id="f-001", + status=FileStatus.CHUNKING, + message="all good", + took=0.42, + ) + + @pytest.mark.asyncio + async def test_send_report_with_socket_calls_send_json(self): + """send_report must call socket.send_json with the correct payload.""" + mock_socket = MagicMock() + mock_socket.send_json = AsyncMock() + lm = LoggerManager(socket=mock_socket) + + await lm.send_report( + file_Id="f-001", + status=FileStatus.DONE, + message="finished", + took=1.23, + ) + + mock_socket.send_json.assert_awaited_once() + sent_payload = mock_socket.send_json.call_args[0][0] + assert sent_payload["fileID"] == "f-001" + assert sent_payload["status"] == FileStatus.DONE + assert sent_payload["message"] == "finished" + assert sent_payload["took"] == 1.23 + + @pytest.mark.asyncio + async def test_send_report_without_socket_skips_send_json(self): + """Without a socket, send_json must never be called.""" + mock_socket = MagicMock() + mock_socket.send_json = AsyncMock() + + # Deliberately do NOT pass the socket + lm = LoggerManager() + await lm.send_report("f-002", FileStatus.ERROR, "oops", 0.0) + + mock_socket.send_json.assert_not_awaited() + + +class TestLoggerManagerCreateNewDocument: + @pytest.mark.asyncio + async def test_create_new_document_without_socket_does_not_crash(self): + """create_new_document must not raise when socket is None.""" + lm = LoggerManager() + await lm.create_new_document( + new_file_id="new-001", + document_name="renamed.txt", + original_file_id="orig-001", + ) + + @pytest.mark.asyncio + async def test_create_new_document_with_socket_sends_correct_payload(self): + """create_new_document must send all three identifiers via socket.""" + mock_socket = MagicMock() + mock_socket.send_json = AsyncMock() + lm = LoggerManager(socket=mock_socket) + + await lm.create_new_document( + new_file_id="new-001", + document_name="renamed.txt", + original_file_id="orig-001", + ) + + mock_socket.send_json.assert_awaited_once() + payload = mock_socket.send_json.call_args[0][0] + assert payload["new_file_id"] == "new-001" + assert payload["filename"] == "renamed.txt" + assert payload["original_file_id"] == "orig-001" diff --git a/goldenverba/verba_manager.py b/goldenverba/verba_manager.py deleted file mode 100644 index 16412fb5ae..0000000000 --- a/goldenverba/verba_manager.py +++ /dev/null @@ -1,827 +0,0 @@ -import os -import importlib -import math -import json -from datetime import datetime - -from dotenv import load_dotenv -from wasabi import msg -import asyncio - -from copy import deepcopy -import hashlib - -from goldenverba.server.helpers import LoggerManager -from weaviate.client import WeaviateAsyncClient - -from goldenverba.components.document import Document -from goldenverba.server.types import ( - FileConfig, - FileStatus, - ChunkScore, - Credentials, -) - -from goldenverba.components.managers import ( - ReaderManager, - ChunkerManager, - EmbeddingManager, - RetrieverManager, - GeneratorManager, - WeaviateManager, -) - -load_dotenv() - - -class VerbaManager: - """Manages all Verba Components.""" - - def __init__(self) -> None: - self.reader_manager = ReaderManager() - self.chunker_manager = ChunkerManager() - self.embedder_manager = EmbeddingManager() - self.retriever_manager = RetrieverManager() - self.generator_manager = GeneratorManager() - self.weaviate_manager = WeaviateManager() - self.rag_config_uuid = "e0adcc12-9bad-4588-8a1e-bab0af6ed485" - self.theme_config_uuid = "baab38a7-cb51-4108-acd8-6edeca222820" - self.user_config_uuid = "f53f7738-08be-4d5a-b003-13eb4bf03ac7" - self.environment_variables = {} - self.installed_libraries = {} - - self.verify_installed_libraries() - self.verify_variables() - - async def connect(self, credentials: Credentials, port: str = "8080"): - start_time = asyncio.get_event_loop().time() - try: - client = await self.weaviate_manager.connect( - credentials.deployment, credentials.url, credentials.key, port - ) - except Exception as e: - raise e - if client: - initialized = await self.weaviate_manager.verify_collection( - client, self.weaviate_manager.config_collection_name - ) - if initialized: - end_time = asyncio.get_event_loop().time() - msg.info(f"Connection time: {end_time - start_time:.2f} seconds") - return client - - async def disconnect(self, client): - start_time = asyncio.get_event_loop().time() - result = await self.weaviate_manager.disconnect(client) - end_time = asyncio.get_event_loop().time() - msg.info(f"Disconnection time: {end_time - start_time:.2f} seconds") - return result - - async def get_deployments(self): - deployments = { - "WEAVIATE_URL_VERBA": ( - os.getenv("WEAVIATE_URL_VERBA") - if os.getenv("WEAVIATE_URL_VERBA") - else "" - ), - "WEAVIATE_API_KEY_VERBA": ( - os.getenv("WEAVIATE_API_KEY_VERBA") - if os.getenv("WEAVIATE_API_KEY_VERBA") - else "" - ), - } - return deployments - - # Import - - async def import_document( - self, client, fileConfig: FileConfig, logger: LoggerManager = LoggerManager() - ): - try: - loop = asyncio.get_running_loop() - start_time = loop.time() - - duplicate_uuid = await self.weaviate_manager.exist_document_name( - client, fileConfig.filename - ) - if duplicate_uuid is not None and not fileConfig.overwrite: - raise Exception(f"{fileConfig.filename} already exists in Verba") - elif duplicate_uuid is not None and fileConfig.overwrite: - await self.weaviate_manager.delete_document(client, duplicate_uuid) - await logger.send_report( - fileConfig.fileID, - status=FileStatus.STARTING, - message=f"Overwriting {fileConfig.filename}", - took=0, - ) - else: - await logger.send_report( - fileConfig.fileID, - status=FileStatus.STARTING, - message="Starting Import", - took=0, - ) - - documents = await self.reader_manager.load( - fileConfig.rag_config["Reader"].selected, fileConfig, logger - ) - - tasks = [ - self.process_single_document(client, doc, fileConfig, logger) - for doc in documents - ] - - results = await asyncio.gather(*tasks, return_exceptions=True) - successful_tasks = sum( - 1 for result in results if not isinstance(result, Exception) - ) - - if successful_tasks > 1: - await logger.send_report( - fileConfig.fileID, - status=FileStatus.INGESTING, - message=f"Imported {fileConfig.filename} and it's {successful_tasks} documents into Weaviate", - took=round(loop.time() - start_time, 2), - ) - elif successful_tasks == 1: - await logger.send_report( - fileConfig.fileID, - status=FileStatus.INGESTING, - message=f"Imported {fileConfig.filename} and {len(documents[0].chunks)} chunks into Weaviate", - took=round(loop.time() - start_time, 2), - ) - elif ( - successful_tasks == 0 - and len(results) == 1 - and isinstance(results[0], Exception) - ): - msg.fail( - f"No documents imported {successful_tasks} of {len(results)} succesful tasks" - ) - raise results[0] - else: - raise Exception( - f"No documents imported {successful_tasks} of {len(results)} succesful tasks" - ) - - await logger.send_report( - fileConfig.fileID, - status=FileStatus.DONE, - message=f"Import for {fileConfig.filename} completed successfully", - took=round(loop.time() - start_time, 2), - ) - - except Exception as e: - await logger.send_report( - fileConfig.fileID, - status=FileStatus.ERROR, - message=f"Import for {fileConfig.filename} failed: {str(e)}", - took=0, - ) - return - - async def process_single_document( - self, - client, - document: Document, - fileConfig: FileConfig, - logger: LoggerManager, - ): - loop = asyncio.get_running_loop() - start_time = loop.time() - - if fileConfig.isURL: - currentFileConfig = deepcopy(fileConfig) - currentFileConfig.fileID = fileConfig.fileID + document.title - currentFileConfig.isURL = False - currentFileConfig.filename = document.title - await logger.create_new_document( - fileConfig.fileID + document.title, - document.title, - fileConfig.fileID, - ) - else: - currentFileConfig = fileConfig - - try: - duplicate_uuid = await self.weaviate_manager.exist_document_name( - client, document.title - ) - if duplicate_uuid is not None and not currentFileConfig.overwrite: - raise Exception(f"{document.title} already exists in Verba") - elif duplicate_uuid is not None and currentFileConfig.overwrite: - await self.weaviate_manager.delete_document(client, duplicate_uuid) - - chunk_task = asyncio.create_task( - self.chunker_manager.chunk( - currentFileConfig.rag_config["Chunker"].selected, - currentFileConfig, - [document], - self.embedder_manager.embedders[ - currentFileConfig.rag_config["Embedder"].selected - ], - logger, - ) - ) - chunked_documents = await chunk_task - - embedding_task = asyncio.create_task( - self.embedder_manager.vectorize( - currentFileConfig.rag_config["Embedder"].selected, - currentFileConfig, - chunked_documents, - logger, - ) - ) - vectorized_documents = await embedding_task - - for document in vectorized_documents: - ingesting_task = asyncio.create_task( - self.weaviate_manager.import_document( - client, - document, - currentFileConfig.rag_config["Embedder"] - .components[fileConfig.rag_config["Embedder"].selected] - .config["Model"] - .value, - ) - ) - await ingesting_task - - await logger.send_report( - currentFileConfig.fileID, - status=FileStatus.INGESTING, - message=f"Imported {currentFileConfig.filename} into Weaviate", - took=round(loop.time() - start_time, 2), - ) - - await logger.send_report( - currentFileConfig.fileID, - status=FileStatus.DONE, - message=f"Import for {currentFileConfig.filename} completed successfully", - took=round(loop.time() - start_time, 2), - ) - except Exception as e: - await logger.send_report( - currentFileConfig.fileID, - status=FileStatus.ERROR, - message=f"Import for {fileConfig.filename} failed: {str(e)}", - took=round(loop.time() - start_time, 2), - ) - raise Exception(f"Import for {fileConfig.filename} failed: {str(e)}") - - # Configuration - - def create_config(self) -> dict: - """Creates the RAG Configuration and returns the full Verba Config with also Settings""" - - available_environments = self.environment_variables - available_libraries = self.installed_libraries - - readers = self.reader_manager.readers - reader_config = { - "components": { - reader: readers[reader].get_meta( - available_environments, available_libraries - ) - for reader in readers - }, - "selected": list(readers.values())[0].name, - } - - chunkers = self.chunker_manager.chunkers - chunkers_config = { - "components": { - chunker: chunkers[chunker].get_meta( - available_environments, available_libraries - ) - for chunker in chunkers - }, - "selected": list(chunkers.values())[0].name, - } - - embedders = self.embedder_manager.embedders - embedder_config = { - "components": { - embedder: embedders[embedder].get_meta( - available_environments, available_libraries - ) - for embedder in embedders - }, - "selected": list(embedders.values())[0].name, - } - - retrievers = self.retriever_manager.retrievers - retrievers_config = { - "components": { - retriever: retrievers[retriever].get_meta( - available_environments, available_libraries - ) - for retriever in retrievers - }, - "selected": list(retrievers.values())[0].name, - } - - generators = self.generator_manager.generators - generator_config = { - "components": { - generator: generators[generator].get_meta( - available_environments, available_libraries - ) - for generator in generators - }, - "selected": list(generators.values())[0].name, - } - - return { - "Reader": reader_config, - "Chunker": chunkers_config, - "Embedder": embedder_config, - "Retriever": retrievers_config, - "Generator": generator_config, - } - - def create_user_config(self) -> dict: - return {"getting_started": False} - - async def set_theme_config(self, client, config: dict): - await self.weaviate_manager.set_config(client, self.theme_config_uuid, config) - - async def set_rag_config(self, client, config: dict): - await self.weaviate_manager.set_config(client, self.rag_config_uuid, config) - - async def set_user_config(self, client, config: dict): - await self.weaviate_manager.set_config(client, self.user_config_uuid, config) - - async def load_rag_config(self, client): - """Check if a Configuration File exists in the database, if yes, check if corrupted. Returns a valid configuration file""" - loaded_config = await self.weaviate_manager.get_config( - client, self.rag_config_uuid - ) - new_config = self.create_config() - if loaded_config is not None: - if self.verify_config(loaded_config, new_config): - msg.info("Using Existing RAG Configuration") - return loaded_config - else: - msg.info("Using New RAG Configuration") - await self.set_rag_config(client, new_config) - return new_config - else: - msg.info("Using New RAG Configuration") - return new_config - - async def load_theme_config(self, client): - loaded_config = await self.weaviate_manager.get_config( - client, self.theme_config_uuid - ) - - if loaded_config is None: - return None, None - - return loaded_config["theme"], loaded_config["themes"] - - async def load_user_config(self, client): - loaded_config = await self.weaviate_manager.get_config( - client, self.user_config_uuid - ) - - if loaded_config is None: - return self.create_user_config() - - return loaded_config - - def verify_config(self, a: dict, b: dict) -> bool: - # Check Settings ( RAG & Settings ) - try: - if os.getenv("VERBA_PRODUCTION") == "Demo": - return True - for a_component_key, b_component_key in zip(a, b): - if a_component_key != b_component_key: - msg.fail( - f"Config Validation Failed, component name mismatch: {a_component_key} != {b_component_key}" - ) - return False - - a_component = a[a_component_key]["components"] - b_component = b[b_component_key]["components"] - - if len(a_component) != len(b_component): - msg.fail( - f"Config Validation Failed, {a_component_key} component count mismatch: {len(a_component)} != {len(b_component)}" - ) - return False - - for a_rag_component_key, b_rag_component_key in zip( - a_component, b_component - ): - if a_rag_component_key != b_rag_component_key: - msg.fail( - f"Config Validation Failed, component name mismatch: {a_rag_component_key} != {b_rag_component_key}" - ) - return False - a_rag_component = a_component[a_rag_component_key] - b_rag_component = b_component[b_rag_component_key] - - a_config = a_rag_component["config"] - b_config = b_rag_component["config"] - - if len(a_config) != len(b_config): - msg.fail( - f"Config Validation Failed, component config count mismatch: {len(a_config)} != {len(b_config)}" - ) - return False - - for a_config_key, b_config_key in zip(a_config, b_config): - if a_config_key != b_config_key: - msg.fail( - f"Config Validation Failed, component name mismatch: {a_config_key} != {b_config_key}" - ) - return False - - a_setting = a_config[a_config_key] - b_setting = b_config[b_config_key] - - if a_setting["description"] != b_setting["description"]: - msg.fail( - f"Config Validation Failed, description mismatch: {a_setting['description']} != {b_setting['description']}" - ) - return False - - if sorted(a_setting["values"]) != sorted(b_setting["values"]): - msg.fail( - f"Config Validation Failed, values mismatch: {a_setting['values']} != {b_setting['values']}" - ) - return False - - return True - - except Exception as e: - msg.fail(f"Config Validation failed: {str(e)}") - return False - - async def reset_rag_config(self, client): - msg.info("Resetting RAG Configuration") - await self.weaviate_manager.reset_config(client, self.rag_config_uuid) - - async def reset_theme_config(self, client): - msg.info("Resetting Theme Configuration") - await self.weaviate_manager.reset_config(client, self.theme_config_uuid) - - async def reset_user_config(self, client): - msg.info("Resetting User Configuration") - await self.weaviate_manager.reset_config(client, self.user_config_uuid) - - # Environment and Libraries - - def verify_installed_libraries(self) -> None: - """ - Checks which libraries are installed and fills out the self.installed_libraries dictionary for the frontend to access, this will be displayed in the status page. - """ - reader = [ - lib - for reader in self.reader_manager.readers - for lib in self.reader_manager.readers[reader].requires_library - ] - chunker = [ - lib - for chunker in self.chunker_manager.chunkers - for lib in self.chunker_manager.chunkers[chunker].requires_library - ] - embedder = [ - lib - for embedder in self.embedder_manager.embedders - for lib in self.embedder_manager.embedders[embedder].requires_library - ] - retriever = [ - lib - for retriever in self.retriever_manager.retrievers - for lib in self.retriever_manager.retrievers[retriever].requires_library - ] - generator = [ - lib - for generator in self.generator_manager.generators - for lib in self.generator_manager.generators[generator].requires_library - ] - - required_libraries = reader + chunker + embedder + retriever + generator - unique_libraries = set(required_libraries) - - for lib in unique_libraries: - try: - importlib.import_module(lib) - self.installed_libraries[lib] = True - except Exception: - self.installed_libraries[lib] = False - - def verify_variables(self) -> None: - """ - Checks which environment variables are installed and fills out the self.environment_variables dictionary for the frontend to access. - """ - reader = [ - lib - for reader in self.reader_manager.readers - for lib in self.reader_manager.readers[reader].requires_env - ] - chunker = [ - lib - for chunker in self.chunker_manager.chunkers - for lib in self.chunker_manager.chunkers[chunker].requires_env - ] - embedder = [ - lib - for embedder in self.embedder_manager.embedders - for lib in self.embedder_manager.embedders[embedder].requires_env - ] - retriever = [ - lib - for retriever in self.retriever_manager.retrievers - for lib in self.retriever_manager.retrievers[retriever].requires_env - ] - generator = [ - lib - for generator in self.generator_manager.generators - for lib in self.generator_manager.generators[generator].requires_env - ] - - required_envs = reader + chunker + embedder + retriever + generator - unique_envs = set(required_envs) - - for env in unique_envs: - if os.environ.get(env) is not None: - self.environment_variables[env] = True - else: - self.environment_variables[env] = False - - # Document Content Retrieval - - async def get_content( - self, - client, - uuid: str, - page: int, - chunkScores: list[ChunkScore], - ): - chunks_per_page = 10 - content_pieces = [] - total_batches = 0 - - # Return Chunks with surrounding context - if len(chunkScores) > 0: - if page > len(chunkScores): - page = 0 - - total_batches = len(chunkScores) - chunk = await self.weaviate_manager.get_chunk( - client, chunkScores[page].uuid, chunkScores[page].embedder - ) - - before_ids = [ - i - for i in range( - max(0, chunkScores[page].chunk_id - int(chunks_per_page / 2)), - chunkScores[page].chunk_id, - ) - ] - if before_ids: - chunks_before_chunk = await self.weaviate_manager.get_chunk_by_ids( - client, - chunkScores[page].embedder, - uuid, - ids=[ - i - for i in range( - max( - 0, chunkScores[page].chunk_id - int(chunks_per_page / 2) - ), - chunkScores[page].chunk_id, - ) - ], - ) - before_content = "".join( - [ - chunk.properties["content_without_overlap"] - for chunk in chunks_before_chunk - ] - ) - else: - before_content = "" - - after_ids = [ - i - for i in range( - chunkScores[page].chunk_id + 1, - chunkScores[page].chunk_id + int(chunks_per_page / 2), - ) - ] - if after_ids: - chunks_after_chunk = await self.weaviate_manager.get_chunk_by_ids( - client, - chunkScores[page].embedder, - uuid, - ids=[ - i - for i in range( - chunkScores[page].chunk_id + 1, - chunkScores[page].chunk_id + int(chunks_per_page / 2), - ) - ], - ) - after_content = "".join( - [ - chunk.properties["content_without_overlap"] - for chunk in chunks_after_chunk - ] - ) - else: - after_content = "" - - content_pieces.append( - { - "content": before_content, - "chunk_id": 0, - "score": 0, - "type": "text", - } - ) - content_pieces.append( - { - "content": chunk["content_without_overlap"], - "chunk_id": chunkScores[page].chunk_id, - "score": chunkScores[page].score, - "type": "extract", - } - ) - content_pieces.append( - { - "content": after_content, - "chunk_id": 0, - "score": 0, - "type": "text", - } - ) - - # Return Content based on Page - else: - document = await self.weaviate_manager.get_document( - client, uuid, properties=["meta"] - ) - config = json.loads(document["meta"]) - embedder = config["Embedder"]["config"]["Model"]["value"] - request_chunk_ids = [ - i - for i in range( - chunks_per_page * (page + 1) - chunks_per_page, - chunks_per_page * (page + 1), - ) - ] - - chunks = await self.weaviate_manager.get_chunk_by_ids( - client, embedder, uuid, request_chunk_ids - ) - - total_chunks = await self.weaviate_manager.get_chunk_count( - client, embedder, uuid - ) - total_batches = int(math.ceil(total_chunks / chunks_per_page)) - - content = "".join( - [chunk.properties["content_without_overlap"] for chunk in chunks] - ) - - content_pieces.append( - { - "content": content, - "chunk_id": 0, - "score": 0, - "type": "text", - } - ) - - return (content_pieces, total_batches) - - # Retrieval Augmented Generation - - async def retrieve_chunks( - self, - client, - query: str, - rag_config: dict, - labels: list[str] = [], - document_uuids: list[str] = [], - ): - retriever = rag_config["Retriever"].selected - embedder = rag_config["Embedder"].selected - - await self.weaviate_manager.add_suggestion(client, query) - - vector = await self.embedder_manager.vectorize_query( - embedder, query, rag_config - ) - documents, context = await self.retriever_manager.retrieve( - client, - retriever, - query, - vector, - rag_config, - self.weaviate_manager, - labels, - document_uuids, - ) - - return (documents, context) - - async def generate_stream_answer( - self, - rag_config: dict, - query: str, - context: str, - conversation: list[dict], - ): - - full_text = "" - async for result in self.generator_manager.generate_stream( - rag_config, query, context, conversation - ): - full_text += result["message"] - yield result - - -class ClientManager: - def __init__(self) -> None: - self.clients: dict[str, dict] = {} - self.manager: VerbaManager = VerbaManager() - self.max_time: int = 10 - self.locks: dict[str, asyncio.Lock] = {} - - def hash_credentials(self, credentials: Credentials) -> str: - cred_string = f"{credentials.deployment}:{credentials.url}:{credentials.key}" - return hashlib.sha256(cred_string.encode()).hexdigest() - - def get_or_create_lock(self, cred_hash: str) -> asyncio.Lock: - if cred_hash not in self.locks: - self.locks[cred_hash] = asyncio.Lock() - return self.locks[cred_hash] - - def heartbeat(self): - msg.info(f"{len(self.clients)} clients connected") - for cred_hash, client in self.clients.items(): - msg.info(f"Client {cred_hash} connected at {client['timestamp']}") - - async def connect( - self, credentials: Credentials, port: str = "8080" - ) -> WeaviateAsyncClient: - - _credentials = credentials - - if not _credentials.url and not _credentials.key: - _credentials.url = os.environ.get("WEAVIATE_URL_VERBA", "") - _credentials.key = os.environ.get("WEAVIATE_API_KEY_VERBA", "") - - cred_hash = self.hash_credentials(_credentials) - - lock = self.get_or_create_lock(cred_hash) - async with lock: - if cred_hash in self.clients: - msg.info("Found existing Client") - return self.clients[cred_hash]["client"] - else: - msg.warn("Connecting new Client") - try: - client = await self.manager.connect(_credentials, port) - if client: - self.clients[cred_hash] = { - "client": client, - "timestamp": datetime.now(), - } - return client - else: - raise Exception("Client not created") - except Exception as e: - raise e - - async def disconnect(self): - msg.warn("Disconnecting Clients!") - for cred_hash, client in self.clients.items(): - await self.manager.disconnect(client["client"]) - - async def clean_up(self): - msg.info("Cleaning Clients Cache") - current_time = datetime.now() - clients_to_remove = [] - - for cred_hash, client_data in self.clients.items(): - time_difference = current_time - client_data["timestamp"] - if time_difference.total_seconds() / 60 > self.max_time: - clients_to_remove.append(cred_hash) - client: WeaviateAsyncClient = client_data["client"] - if not await client.is_ready(): - clients_to_remove.append(cred_hash) - - for cred_hash in clients_to_remove: - await self.manager.disconnect(self.clients[cred_hash]["client"]) - del self.clients[cred_hash] - msg.warn(f"Removed client: {cred_hash}") - - msg.info(f"Cleaned up {len(clients_to_remove)} clients") - self.heartbeat() diff --git a/pypi_commands.sh b/pypi_commands.sh deleted file mode 100644 index 6c27af6d3e..0000000000 --- a/pypi_commands.sh +++ /dev/null @@ -1,3 +0,0 @@ -python setup.py sdist bdist_wheel - -twine upload dist/* \ No newline at end of file diff --git a/ruff.toml b/ruff.toml new file mode 100644 index 0000000000..c367ad75d0 --- /dev/null +++ b/ruff.toml @@ -0,0 +1,24 @@ +target-version = "py310" +line-length = 100 + +[lint] +select = [ + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort + "UP", # pyupgrade + "B", # flake8-bugbear +] +ignore = [ + "E501", # line too long (handled by formatter) + "B008", # do not perform function calls in argument defaults (common in FastAPI) + "B904", # raise from within except (can be addressed incrementally) +] + +[lint.isort] +known-first-party = ["goldenverba"] + +[format] +quote-style = "double" +indent-style = "space" diff --git a/setup.py b/setup.py index 9f20bc4c10..81e1586ac3 100644 --- a/setup.py +++ b/setup.py @@ -24,7 +24,7 @@ ], include_package_data=True, install_requires=[ - "weaviate-client==4.9.6", + "weaviate-client>=4.20.5", "python-dotenv==1.0.0", "openpyxl==3.1.5", "wasabi==1.1.2", @@ -33,7 +33,6 @@ "gunicorn==22.0.0", "click==8.1.7", "xlrd==2.0.2", - "asyncio==3.4.3", "tiktoken==0.6.0", "requests==2.31.0", "pypdf==4.3.1", @@ -42,9 +41,10 @@ "langchain-text-splitters==0.2.2", "spacy==3.7.5", "aiohttp==3.9.5", + "httpx>=0.27.0", "markdownify==0.13.1", "aiofiles==24.1.0", - "assemblyai==0.33.0", + "faster-whisper>=1.0.0", "beautifulsoup4==4.12.3", "langdetect==1.0.9", ],