From ea2fd2d23af1f017e6bd7e67163e11a212c4abb9 Mon Sep 17 00:00:00 2001 From: Matuck Date: Mon, 7 Sep 2026 19:16:36 -0300 Subject: [PATCH 1/4] Rewrite README, add MIT LICENSE and .env.example README describes only what the code does today: section chunking, OpenAI embeddings cached in JSON, in-memory dot-product ranking, gpt-4o-mini answers, conversation history in PostgreSQL via SQLAlchemy and Alembic. The curl example is marked as illustrative. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QVi7kKzGQhgPKX8sDzTu6i --- .env.example | 2 + LICENSE | 21 +++ README.md | 452 +++++++++------------------------------------------ 3 files changed, 101 insertions(+), 374 deletions(-) create mode 100644 .env.example create mode 100644 LICENSE diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..cb60584 --- /dev/null +++ b/.env.example @@ -0,0 +1,2 @@ +OPENAI_API_KEY= +DATABASE_URL= diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..e6c944f --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Lucca Matuck + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index c3694d3..9de10a4 100644 --- a/README.md +++ b/README.md @@ -1,412 +1,116 @@ -# RAG Chatbot com Histórico de Conversas +A retrieval-augmented generation API: a FastAPI service that splits a Markdown knowledge base into sections, embeds them with the OpenAI embeddings API, ranks them in memory by dot-product similarity, answers questions with `gpt-4o-mini` grounded on the three best-matching chunks, and stores each conversation in PostgreSQL. -API em FastAPI conectada a um modelo da OpenAI com armazenamento de histórico em PostgreSQL. +## Architecture -## ✨ Funcionalidades - -- **RAG (Retrieval-Augmented Generation)**: Busca informações relevantes na base de conhecimento -- **Histórico de Conversas**: Armazena e recupera conversas completas usando PostgreSQL -- **Sessões de Chat**: Suporte a múltiplas sessões de conversa independentes -- **API REST**: Endpoints para perguntas, histórico e saúde do sistema - -## 📋 Requisitos - -- macOS, Linux ou Windows -- Python 3.11+ -- PostgreSQL 12+ -- VS Code opcional -- Chave da OpenAI - -## 🚀 Instalação Rápida - -### 1. Clonar/Configurar Ambiente - -```bash -cd ~/projects/rag-chatbot -python3 -m venv .venv -source .venv/bin/activate -pip install -r requirements.txt ``` - -### 2. Configurar PostgreSQL - -**macOS:** -```bash -brew install postgresql -brew services start postgresql -createdb rag_chatbot +client + | POST /ask {question, session_id?, history?} + v +app/main.py (FastAPI: GET /health, POST /ask, GET /conversations/{session_id}) + | + |-- startup: data/knowledge_base.txt + | -> app/knowledge.py chunk_by_sections() (split on "# " / "## " headings) + | -> app/embedding.py OpenAI text-embedding-3-small, one call per chunk + | -> data/chunk_embeddings.json (cache, reused while the file's mtime is + | unchanged, git-ignored) + | + |-- app/services/rag_service.py embed the question, dot product against every cached + | chunk, keep the top 3 as context + "sources" + |-- app/llm.py OpenAI chat completion (gpt-4o-mini): + | system prompt + context + prior turns + question + |-- app/services/database_service.py SQLAlchemy engine from DATABASE_URL, + | reads/writes conversations and messages + v +PostgreSQL (tables: conversations, messages; schema in alembic/versions/001_*.py) ``` -**Ubuntu/Debian:** -```bash -sudo apt update -sudo apt install postgresql postgresql-contrib -sudo systemctl start postgresql -sudo -u postgres createdb rag_chatbot -``` - -### 3. Configurar Variáveis de Ambiente +Components, by file: -Edite o arquivo `.env`: - -```env -OPENAI_API_KEY=sk-sua-chave-aqui -DATABASE_URL=postgresql://localhost/rag_chatbot -``` - -### 4. Inicializar Banco de Dados - -```bash -python init_db.py -``` +- `app/main.py`: builds the FastAPI app and, at import time, loads and embeds the knowledge base. `POST /ask` takes a `question`, an optional `session_id` (a UUID4 is generated if absent) and an optional `history` list. If `history` is not sent, prior messages for that session are loaded from the database. After the LLM answers, the user message and the assistant message are saved. `GET /conversations/{session_id}` returns the stored messages. `GET /health` returns `{"status": "ok"}`. +- `app/knowledge.py`: `load_system_prompt` reads `prompts/system_prompt.txt`; `chunk_by_sections` starts a new chunk at every line beginning with `# ` or `## ` (`### ` headings stay inside their parent chunk). The sample knowledge base yields 14 chunks. `load_knowledge_base` and `chunk_text` in the same file are not called anywhere. +- `app/embedding.py`: reads `data/knowledge_base.txt`, calls `text-embedding-3-small`, computes the dot product in plain Python, and caches chunk embeddings in `data/chunk_embeddings.json`. The cache is reused only if the knowledge file's modification time is unchanged; a change to the chunking function alone does not invalidate it. +- `app/services/rag_service.py`: retrieval. Returns the concatenated top-3 chunks as context and a `sources` list with each chunk's score and its first 200 characters. +- `app/llm.py`: sends two system messages (the prompt file, then `Contexto:\n`), then the history turns, then the question, to `gpt-4o-mini`. +- `app/models.py`: the Pydantic request models (`QuestionRequest`, `Message`) and the SQLAlchemy ORM models (`Conversation`, `MessageDB`) live in the same file. +- `app/services/database_service.py`: `DatabaseService`, a module-level singleton created at import. It opens a session per operation. The API uses two of its methods, `save_message_to_conversation` and `get_conversation_history`; `create_conversation`, `get_conversation` and `add_message` are defined but not called. +- `alembic/`, `alembic.ini`: one migration (`001`) creating `conversations` and `messages`. `alembic/env.py` reads the connection string from `alembic.ini`, not from `.env`. +- `init_db.py`: alternative to Alembic; runs `Base.metadata.create_all` with `DATABASE_URL` from `.env` and does not record a migration version. +- `prompts/system_prompt.txt`: the assistant persona (a support agent for a fictional company "Acme"), in Brazilian Portuguese, instructing the model to answer only from the supplied context. +- `data/knowledge_base.txt`: the sample knowledge base (plans, billing, support policies of "Acme"), in Portuguese, about 300 lines. -### 5. Executar +Retrieval happens in process, against the list of chunk embeddings held in memory since startup. PostgreSQL stores chat history only; there is no vector extension and no vector index. -```bash -./run.sh -``` - -Ou diretamente: -```bash -source .venv/bin/activate -uvicorn app.main:app --reload -``` - -## Requisitos - -- macOS, Linux ou Windows -- Python 3.11+ -- PostgreSQL database -- VS Code opcional -- uma chave da OpenAI - -## Estrutura esperada - -```text -rag-chatbot/ - app/ - main.py - .venv/ - .env - .gitignore - requirements.txt - run.sh -``` - -## Instalação - -### 1. Entrar na pasta do projeto - -```bash -cd ~/projects/rag-chatbot -``` - -### 2. Criar o ambiente virtual - -```bash -python3 -m venv .venv -source .venv/bin/activate -``` +## Example -### 3. Instalar dependências +The example below is illustrative, not captured from a live run (the app needs a valid OpenAI key even to start: the client is constructed at import, and without a cache the embeddings are generated at import). The request and response fields match the code; the answer text and scores are placeholders. The previews are the first characters of real chunks of `data/knowledge_base.txt`. ```bash -pip install -r requirements.txt -``` - -## Configurar a chave da OpenAI - -Crie um arquivo chamado `.env` na raiz do projeto com este formato: - -```env -OPENAI_API_KEY=sk-sua-chave-aqui -DATABASE_URL=postgresql://user:password@localhost/rag_chatbot +curl -s -X POST http://127.0.0.1:8000/ask \ + -H "Content-Type: application/json" \ + -d '{"question": "Quantos usuários o Plano Basic permite?"}' ``` -## 📡 API Endpoints - -### GET `/health` -Verifica se o servidor está funcionando. - -**Resposta:** -```json -{"status": "ok"} -``` - -### POST `/ask` -Envia uma pergunta e recebe resposta com contexto. - -**Request Body:** ```json { - "question": "Qual é a capital do Brasil?", - "session_id": "abc-123-def", - "history": [ - {"role": "user", "content": "Olá"}, - {"role": "assistant", "content": "Olá! Como posso ajudar?"} - ] + "answer": "O Plano Basic permite até 5 usuários. ...", + "sources": [ + {"score": 0.61, "text_preview": "## Usuários e limites\n\n### Limites por plano\n- Basic: até 5 usuários\n- Pro: até ..."}, + {"score": 0.55, "text_preview": "## Planos disponíveis\n\n### Plano Basic\nO Plano Basic é indicado para profissiona..."}, + {"score": 0.47, "text_preview": "## Comparação rápida entre planos\n\nResumo:\n- Basic: solução inicial para uso sim..."} + ], + "session_id": "8f3c2c0e-1b7e-4b2e-9a4e-2f1c3d4e5f60" } ``` -**Resposta:** -```json -{ - "answer": "A capital do Brasil é Brasília.", - "sources": ["fonte1", "fonte2"], - "session_id": "abc-123-def" -} -``` - -### GET `/conversations/{session_id}` -Recupera histórico completo de uma conversa. - -**Resposta:** -```json -{ - "session_id": "abc-123-def", - "messages": [ - {"role": "user", "content": "Olá"}, - {"role": "assistant", "content": "Olá! Como posso ajudar?"}, - {"role": "user", "content": "Qual é a capital do Brasil?"}, - {"role": "assistant", "content": "A capital do Brasil é Brasília."} - ] -} -``` - - -## Exemplo de `run.sh` - -Crie um arquivo `run.sh` na raiz: - -```bash -#!/bin/bash -cd "$(dirname "$0")" -source .venv/bin/activate -uvicorn app.main:app --reload -``` - -Deixe executável: - -```bash -chmod +x run.sh -``` - -## Exemplo de `app/main.py` - -```python -from fastapi import FastAPI -from pydantic import BaseModel -from openai import OpenAI -from dotenv import load_dotenv - -load_dotenv() - -app = FastAPI() -client = OpenAI() - - -class QuestionRequest(BaseModel): - question: str - history: Optional[List[Message]] = None - session_id: Optional[str] = None +Sending the returned `session_id` in the next request continues the same conversation. `GET /conversations/` returns `{"session_id": "...", "messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}`. +## Running locally -class Message(BaseModel): - role: str # "user" or "assistant" - content: str +1. Python 3.9 or newer. No version is declared in the repository; 3.9 is the floor implied by the built-in generic type hints in the code (such as `list[float]`). Create a virtual environment and install the dependencies; versions are not pinned in `requirements.txt`. - -@app.get("/health") -def health(): - return {"status": "ok"} - - -@app.post("/ask") -def ask(req: QuestionRequest): - response = client.chat.completions.create( - model="gpt-4o-mini", - messages=[ - {"role": "user", "content": req.question} - ] - ) - - answer = response.choices[0].message.content - return {"answer": answer} -``` - -## Rodar o servidor - -Na raiz do projeto: - -```bash -./run.sh -``` - -Ou, sem script: - -```bash -source .venv/bin/activate -uvicorn app.main:app --reload -``` - -## 🧪 Como Testar - -Com o servidor rodando, acesse a documentação interativa: - -```text -http://127.0.0.1:8000/docs -``` - -### Teste Básico - -1. **Health Check:** ```bash - curl http://127.0.0.1:8000/health + python3 -m venv .venv + source .venv/bin/activate + pip install -r requirements.txt ``` -2. **Fazer uma pergunta:** - ```bash - curl -X POST "http://127.0.0.1:8000/ask" \ - -H "Content-Type: application/json" \ - -d '{"question": "Olá, como você funciona?"}' - ``` +2. PostgreSQL. The migration uses only standard column types and no extensions. Create the database: -3. **Ver histórico:** ```bash - curl "http://127.0.0.1:8000/conversations/$(uuidgen)" + createdb rag_chatbot ``` -## 🔧 Troubleshooting - -### Erro de conexão com PostgreSQL -- Verifique se PostgreSQL está rodando: `brew services list` (macOS) ou `sudo systemctl status postgresql` (Linux) -- Confirme que o banco `rag_chatbot` existe: `psql -l` -- Verifique a variável `DATABASE_URL` no `.env` - -### Erro de chave OpenAI -- Confirme que `OPENAI_API_KEY` está definida no `.env` -- Verifique se a chave é válida no [dashboard da OpenAI](https://platform.openai.com/api-keys) - -### ImportError ou ModuleNotFoundError -```bash -source .venv/bin/activate -pip install -r requirements.txt -``` - -## 📁 Estrutura do Projeto +3. Copy `.env.example` to `.env` and fill in both variables. `OPENAI_API_KEY` is read by the `openai` client. `DATABASE_URL` is a SQLAlchemy URL such as `postgresql://USER:PASSWORD@localhost:5432/rag_chatbot`. If `DATABASE_URL` is missing, `database_service.py` falls back to `postgresql://user:password@localhost/rag_chatbot`. -``` -rag-chatbot/ -├── app/ -│ ├── main.py # API FastAPI -│ ├── models.py # Modelos Pydantic e SQLAlchemy -│ ├── services/ -│ │ ├── database_service.py # Serviço de banco de dados -│ │ └── rag_service.py # Serviço RAG -│ ├── embedding.py # Geração de embeddings -│ ├── knowledge.py # Processamento da base de conhecimento -│ └── llm.py # Integração com OpenAI -├── data/ -│ ├── knowledge_base.txt # Base de conhecimento -│ └── chunk_embeddings.json # Embeddings em cache -├── alembic/ # Migrações do banco -├── prompts/ -│ └── system_prompt.txt # Prompt do sistema -├── .env # Variáveis de ambiente -├── requirements.txt # Dependências Python -├── init_db.py # Script de inicialização do banco -└── run.sh # Script de execução -``` - -```json -{"status":"ok"} -``` +4. Create the tables. Alembic takes its URL from the `sqlalchemy.url` line in `alembic.ini` (a placeholder value is committed); set it to the same value as `DATABASE_URL`, then: -### Endpoint de pergunta - -Na página `/docs`: - -1. Abra `POST /ask` -2. Clique em `Try it out` -3. Envie um JSON como este: - -```json -{ - "question": "Explique o que é uma API de forma simples" -} -``` - -Resposta esperada: - -```json -{ - "answer": "..." -} -``` - -## Teste da chave - -Para verificar se o `.env` está sendo lido: - -```bash -python -c "from dotenv import load_dotenv; load_dotenv(); import os; print(os.getenv('OPENAI_API_KEY'))" -``` - -Se aparecer `None`, o `.env` está incorreto ou não está sendo encontrado. - -## Problemas comuns - -### `python-dotenv could not parse statement` - -O arquivo `.env` está mal formatado. - -Formato certo: - -```env -OPENAI_API_KEY=sk-sua-chave-aqui -``` - -### `No module named 'app'` - -Você provavelmente rodou o servidor de dentro da pasta `app`. - -Rode sempre na raiz do projeto: - -```bash -cd ~/projects/rag-chatbot -./run.sh -``` - -### `No API key provided` - -A chave não foi carregada. Verifique: - -- se o `.env` existe -- se `load_dotenv()` está no código -- se a variável está escrita como `OPENAI_API_KEY` - -### `invalid_api_key` + ```bash + alembic upgrade head + ``` -A chave foi lida, mas está errada, incompleta ou expirada. + `python init_db.py` creates the same tables from the ORM models using `.env` instead, without an `alembic_version` table. -## Segurança +5. Start the API. The first start makes one embeddings call per section of the knowledge base and writes `data/chunk_embeddings.json`; later starts reuse it until `data/knowledge_base.txt` changes. -Não suba sua chave para o GitHub. + ```bash + uvicorn app.main:app --reload + ``` -Use um `.gitignore` assim: + `./run.sh` does the same after activating `.venv`. Interactive docs are served at `http://127.0.0.1:8000/docs`. -```gitignore -.venv -.env -__pycache__/ -``` +## What is not here -## Próximos passos +- Automated tests: there is no `tests/` directory. +- Authentication, rate limiting or CORS configuration on the endpoints. +- Deployment: no Dockerfile, container configuration or CI. +- Evaluation of answer quality or retrieval quality. +- A vector store: similarity is computed in memory over a JSON file, which is adequate for the ~300-line sample and not designed for larger corpora. +- Error handling: there is no `except` block in `app/`. If PostgreSQL or the OpenAI API is unreachable, `POST /ask` and `GET /conversations/{session_id}` return a 500. On `/ask`, the database write happens after the LLM call, so a database failure loses an answer that was already generated. +- A limit on conversation length: the full stored history of a session is replayed into every prompt. +- Access control on conversations: anyone who knows a `session_id` can read and extend that conversation. +- Configuration of model names or the number of retrieved chunks; both are hard-coded. +- Accurate in-code comments in two places: the header of `app/main.py` lists `GET /status` and `GET /chunks` endpoints that do not exist, and the docstring of `init_db.py` says it uses Alembic migrations when it calls `create_all`. -Sugestões de evolução: +## Author -- adicionar prompt `system` -- salvar histórico de perguntas e respostas -- criar uma base simples para RAG -- dockerizar a aplicação -- fazer deploy +Lucca Matuck, PhD in Engineering Physics. Signal processing and instrumentation for physical measurement; also builds and operates management systems in production. From d5050134a3d33f6e6ede11206a246fc764d11cde Mon Sep 17 00:00:00 2001 From: Matuck Date: Tue, 8 Sep 2026 13:05:31 -0300 Subject: [PATCH 2/4] Fix stale in-code comments and tighten README wording Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QVi7kKzGQhgPKX8sDzTu6i --- README.md | 5 ++--- app/main.py | 16 ++++++++-------- init_db.py | 3 ++- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 9de10a4..c005d32 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ PostgreSQL (tables: conversations, messages; schema in alembic/versions/001_*.p Components, by file: -- `app/main.py`: builds the FastAPI app and, at import time, loads and embeds the knowledge base. `POST /ask` takes a `question`, an optional `session_id` (a UUID4 is generated if absent) and an optional `history` list. If `history` is not sent, prior messages for that session are loaded from the database. After the LLM answers, the user message and the assistant message are saved. `GET /conversations/{session_id}` returns the stored messages. `GET /health` returns `{"status": "ok"}`. +- `app/main.py`: builds the FastAPI app and, at import time, loads and embeds the knowledge base. `POST /ask` takes a `question`, an optional `session_id` (a UUID4 is generated if absent) and an optional `history` list. If `history` is absent or empty, prior messages for that session are loaded from the database. After the LLM answers, the user message and the assistant message are saved. `GET /conversations/{session_id}` returns the stored messages. `GET /health` returns `{"status": "ok"}`. - `app/knowledge.py`: `load_system_prompt` reads `prompts/system_prompt.txt`; `chunk_by_sections` starts a new chunk at every line beginning with `# ` or `## ` (`### ` headings stay inside their parent chunk). The sample knowledge base yields 14 chunks. `load_knowledge_base` and `chunk_text` in the same file are not called anywhere. - `app/embedding.py`: reads `data/knowledge_base.txt`, calls `text-embedding-3-small`, computes the dot product in plain Python, and caches chunk embeddings in `data/chunk_embeddings.json`. The cache is reused only if the knowledge file's modification time is unchanged; a change to the chunking function alone does not invalidate it. - `app/services/rag_service.py`: retrieval. Returns the concatenated top-3 chunks as context and a `sources` list with each chunk's score and its first 200 characters. @@ -109,8 +109,7 @@ Sending the returned `session_id` in the next request continues the same convers - A limit on conversation length: the full stored history of a session is replayed into every prompt. - Access control on conversations: anyone who knows a `session_id` can read and extend that conversation. - Configuration of model names or the number of retrieved chunks; both are hard-coded. -- Accurate in-code comments in two places: the header of `app/main.py` lists `GET /status` and `GET /chunks` endpoints that do not exist, and the docstring of `init_db.py` says it uses Alembic migrations when it calls `create_all`. ## Author -Lucca Matuck, PhD in Engineering Physics. Signal processing and instrumentation for physical measurement; also builds and operates management systems in production. +Lucca Matuck, PhD in Engineering Physics. [github.com/matuck40](https://github.com/matuck40) diff --git a/app/main.py b/app/main.py index 6f32ffe..dc6f309 100644 --- a/app/main.py +++ b/app/main.py @@ -1,11 +1,11 @@ -## main.py -# FastAPI app that serves as the backend for the RAG chatbot. It loads the knowledge base, creates chunks and embeddings, and provides endpoints for asking questions and checking health/status. - -# The app has the following endpoints: -# - GET /health: Returns a simple health check response. -# - GET /status: Returns the status of the API and the feature being tested. -# - POST /ask: Accepts a question, generates an embedding for it, ranks the knowledge base chunks by similarity, and prints the top 3 most relevant chunks. -# - GET /chunks: Loads the knowledge base, creates chunks and embeddings, and prints their counts for debugging purposes. +# FastAPI app for the RAG chatbot. At import time it loads the knowledge base, +# splits it into chunks and embeds them (cached in data/chunk_embeddings.json). +# +# Endpoints: +# - GET /health: liveness check. +# - POST /ask: embeds the question, retrieves the 3 most similar chunks, asks the +# LLM with the session history, stores both messages and returns the answer. +# - GET /conversations/{session_id}: stored messages of a session. # Import necessary libraries and modules from pathlib import Path diff --git a/init_db.py b/init_db.py index 1a3fa34..13e4286 100644 --- a/init_db.py +++ b/init_db.py @@ -1,7 +1,8 @@ #!/usr/bin/env python3 """ Database initialization script for the RAG chatbot. -This script creates the database tables using Alembic migrations. +Creates the tables directly from the SQLAlchemy models (Base.metadata.create_all), +without recording an Alembic revision. Use `alembic upgrade head` for the versioned path. """ import os From 528bc9c9627d0bb229a3954bd5998f23291bddac Mon Sep 17 00:00:00 2001 From: Matuck Date: Tue, 8 Sep 2026 13:08:20 -0300 Subject: [PATCH 3/4] Clean up the code: lazy OpenAI client, lifespan, schemas split, tests and CI - OpenAI clients are created on first use; embeddings load in the FastAPI lifespan - Pydantic schemas moved to app/schemas.py; app/models.py keeps only the ORM - DatabaseService uses context-managed sessions; unused methods and load_knowledge_base/chunk_text removed - alembic/env.py reads DATABASE_URL (app/__init__.py loads .env); init_db.py removed - .model_dump(), sqlalchemy.orm.declarative_base, timezone-aware UTC helper - tests/ with fakes for OpenAI and the database; ruff config; GitHub Actions workflow - README updated to match Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QVi7kKzGQhgPKX8sDzTu6i --- .github/workflows/tests.yml | 20 ++++++++ .gitignore | 5 +- README.md | 26 +++++----- alembic/env.py | 11 ++-- app/__init__.py | 6 +++ app/embedding.py | 60 +++++++++------------- app/knowledge.py | 20 +------- app/llm.py | 56 +++++++++----------- app/main.py | 68 +++++++++++------------- app/models.py | 30 ++++------- app/schemas.py | 12 +++++ app/services/database_service.py | 88 +++++++++++--------------------- app/services/rag_service.py | 6 ++- init_db.py | 37 -------------- pyproject.toml | 14 +++++ requirements-dev.txt | 4 ++ requirements.txt | 2 +- run.sh | 4 +- tests/test_api.py | 63 +++++++++++++++++++++++ tests/test_embedding.py | 62 ++++++++++++++++++++++ tests/test_knowledge.py | 36 +++++++++++++ 21 files changed, 365 insertions(+), 265 deletions(-) create mode 100644 .github/workflows/tests.yml create mode 100644 app/__init__.py create mode 100644 app/schemas.py delete mode 100644 init_db.py create mode 100644 pyproject.toml create mode 100644 requirements-dev.txt create mode 100644 tests/test_api.py create mode 100644 tests/test_embedding.py create mode 100644 tests/test_knowledge.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..87c8c8e --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,20 @@ +name: tests + +on: + push: + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.11", "3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - run: pip install -r requirements-dev.txt + - run: ruff check . + - run: pytest -q diff --git a/.gitignore b/.gitignore index 77f96f4..90c91fa 100644 --- a/.gitignore +++ b/.gitignore @@ -2,5 +2,6 @@ .env __pycache__ .DS_Store -**/.DS_Store -data/chunk_embeddings.json \ No newline at end of file +data/chunk_embeddings.json +.pytest_cache +.ruff_cache diff --git a/README.md b/README.md index c005d32..d4757a1 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ client v app/main.py (FastAPI: GET /health, POST /ask, GET /conversations/{session_id}) | - |-- startup: data/knowledge_base.txt + |-- startup (lifespan): data/knowledge_base.txt | -> app/knowledge.py chunk_by_sections() (split on "# " / "## " headings) | -> app/embedding.py OpenAI text-embedding-3-small, one call per chunk | -> data/chunk_embeddings.json (cache, reused while the file's mtime is @@ -26,15 +26,15 @@ PostgreSQL (tables: conversations, messages; schema in alembic/versions/001_*.p Components, by file: -- `app/main.py`: builds the FastAPI app and, at import time, loads and embeds the knowledge base. `POST /ask` takes a `question`, an optional `session_id` (a UUID4 is generated if absent) and an optional `history` list. If `history` is absent or empty, prior messages for that session are loaded from the database. After the LLM answers, the user message and the assistant message are saved. `GET /conversations/{session_id}` returns the stored messages. `GET /health` returns `{"status": "ok"}`. -- `app/knowledge.py`: `load_system_prompt` reads `prompts/system_prompt.txt`; `chunk_by_sections` starts a new chunk at every line beginning with `# ` or `## ` (`### ` headings stay inside their parent chunk). The sample knowledge base yields 14 chunks. `load_knowledge_base` and `chunk_text` in the same file are not called anywhere. -- `app/embedding.py`: reads `data/knowledge_base.txt`, calls `text-embedding-3-small`, computes the dot product in plain Python, and caches chunk embeddings in `data/chunk_embeddings.json`. The cache is reused only if the knowledge file's modification time is unchanged; a change to the chunking function alone does not invalidate it. +- `app/main.py`: builds the FastAPI app and, at startup (FastAPI lifespan), loads and embeds the knowledge base. `POST /ask` takes a `question`, an optional `session_id` (a UUID4 is generated if absent) and an optional `history` list. If `history` is absent or empty, prior messages for that session are loaded from the database. After the LLM answers, the user message and the assistant message are saved. `GET /conversations/{session_id}` returns the stored messages. `GET /health` returns `{"status": "ok"}`. +- `app/knowledge.py`: `load_system_prompt` reads `prompts/system_prompt.txt`; `chunk_by_sections` starts a new chunk at every line beginning with `# ` or `## ` (`### ` headings stay inside their parent chunk). The sample knowledge base yields 14 chunks. +- `app/embedding.py`: reads `data/knowledge_base.txt`, calls `text-embedding-3-small` (the OpenAI client is created on first use), computes the dot product in plain Python, and caches chunk embeddings in `data/chunk_embeddings.json`. The cache is reused only if the knowledge file's modification time is unchanged; a change to the chunking function alone does not invalidate it. - `app/services/rag_service.py`: retrieval. Returns the concatenated top-3 chunks as context and a `sources` list with each chunk's score and its first 200 characters. - `app/llm.py`: sends two system messages (the prompt file, then `Contexto:\n`), then the history turns, then the question, to `gpt-4o-mini`. -- `app/models.py`: the Pydantic request models (`QuestionRequest`, `Message`) and the SQLAlchemy ORM models (`Conversation`, `MessageDB`) live in the same file. -- `app/services/database_service.py`: `DatabaseService`, a module-level singleton created at import. It opens a session per operation. The API uses two of its methods, `save_message_to_conversation` and `get_conversation_history`; `create_conversation`, `get_conversation` and `add_message` are defined but not called. -- `alembic/`, `alembic.ini`: one migration (`001`) creating `conversations` and `messages`. `alembic/env.py` reads the connection string from `alembic.ini`, not from `.env`. -- `init_db.py`: alternative to Alembic; runs `Base.metadata.create_all` with `DATABASE_URL` from `.env` and does not record a migration version. +- `app/schemas.py`: the Pydantic request models (`QuestionRequest`, `Message`). `app/models.py`: the SQLAlchemy ORM models (`Conversation`, `MessageDB`). +- `app/services/database_service.py`: `DatabaseService`, a module-level singleton created at import (the engine only connects on first use). It opens a session per operation and has two methods, `save_message_to_conversation` and `get_conversation_history`. +- `alembic/`, `alembic.ini`: one migration (`001`) creating `conversations` and `messages`. `alembic/env.py` uses `DATABASE_URL` when it is set (importing `app` loads `.env`) and falls back to the placeholder in `alembic.ini` otherwise. +- `tests/`: chunking, ranking, cache behaviour and the three endpoints, with the OpenAI calls and the database replaced by fakes. `.github/workflows/tests.yml` runs `ruff` and `pytest` on every push. - `prompts/system_prompt.txt`: the assistant persona (a support agent for a fictional company "Acme"), in Brazilian Portuguese, instructing the model to answer only from the supplied context. - `data/knowledge_base.txt`: the sample knowledge base (plans, billing, support policies of "Acme"), in Portuguese, about 300 lines. @@ -42,7 +42,7 @@ Retrieval happens in process, against the list of chunk embeddings held in memor ## Example -The example below is illustrative, not captured from a live run (the app needs a valid OpenAI key even to start: the client is constructed at import, and without a cache the embeddings are generated at import). The request and response fields match the code; the answer text and scores are placeholders. The previews are the first characters of real chunks of `data/knowledge_base.txt`. +The example below is illustrative, not captured from a live run (without a valid OpenAI key the app cannot finish starting: the embeddings are generated at startup when no cache exists). The request and response fields match the code; the answer text and scores are placeholders. The previews are the first characters of real chunks of `data/knowledge_base.txt`. ```bash curl -s -X POST http://127.0.0.1:8000/ask \ @@ -71,7 +71,7 @@ Sending the returned `session_id` in the next request continues the same convers ```bash python3 -m venv .venv source .venv/bin/activate - pip install -r requirements.txt + pip install -r requirements.txt # or requirements-dev.txt to also get pytest and ruff ``` 2. PostgreSQL. The migration uses only standard column types and no extensions. Create the database: @@ -82,14 +82,12 @@ Sending the returned `session_id` in the next request continues the same convers 3. Copy `.env.example` to `.env` and fill in both variables. `OPENAI_API_KEY` is read by the `openai` client. `DATABASE_URL` is a SQLAlchemy URL such as `postgresql://USER:PASSWORD@localhost:5432/rag_chatbot`. If `DATABASE_URL` is missing, `database_service.py` falls back to `postgresql://user:password@localhost/rag_chatbot`. -4. Create the tables. Alembic takes its URL from the `sqlalchemy.url` line in `alembic.ini` (a placeholder value is committed); set it to the same value as `DATABASE_URL`, then: +4. Create the tables. Alembic reads `DATABASE_URL` from `.env`: ```bash alembic upgrade head ``` - `python init_db.py` creates the same tables from the ORM models using `.env` instead, without an `alembic_version` table. - 5. Start the API. The first start makes one embeddings call per section of the knowledge base and writes `data/chunk_embeddings.json`; later starts reuse it until `data/knowledge_base.txt` changes. ```bash @@ -100,7 +98,7 @@ Sending the returned `session_id` in the next request continues the same convers ## What is not here -- Automated tests: there is no `tests/` directory. +- Tests against a real database or the real OpenAI API: the test suite replaces both with fakes. - Authentication, rate limiting or CORS configuration on the endpoints. - Deployment: no Dockerfile, container configuration or CI. - Evaluation of answer quality or retrieval quality. diff --git a/alembic/env.py b/alembic/env.py index a113dce..8cbae24 100644 --- a/alembic/env.py +++ b/alembic/env.py @@ -1,7 +1,7 @@ +import os from logging.config import fileConfig -from sqlalchemy import engine_from_config -from sqlalchemy import pool +from sqlalchemy import engine_from_config, pool from alembic import context @@ -18,9 +18,14 @@ # for 'autogenerate' support # from myapp import mymodel # target_metadata = mymodel.Base.metadata -from app.models import Base +from app.models import Base # importing `app` also loads .env + target_metadata = Base.metadata +# DATABASE_URL (from the environment or .env) takes precedence over alembic.ini. +if os.getenv("DATABASE_URL"): + config.set_main_option("sqlalchemy.url", os.environ["DATABASE_URL"].replace("%", "%%")) + # other values from the config, defined by the needs of env.py, # can be acquired: # my_important_option = config.get_main_option("my_important_option") diff --git a/app/__init__.py b/app/__init__.py new file mode 100644 index 0000000..96286f6 --- /dev/null +++ b/app/__init__.py @@ -0,0 +1,6 @@ +"""RAG chatbot package. Loading the package reads `.env` so every module +(and Alembic's env.py) sees OPENAI_API_KEY and DATABASE_URL.""" + +from dotenv import load_dotenv + +load_dotenv() diff --git a/app/embedding.py b/app/embedding.py index d289040..f5eaf51 100644 --- a/app/embedding.py +++ b/app/embedding.py @@ -1,53 +1,44 @@ import json +from functools import lru_cache from pathlib import Path + from openai import OpenAI -client = OpenAI() +EMBEDDING_MODEL = "text-embedding-3-small" + + +@lru_cache(maxsize=1) +def get_client() -> OpenAI: + # Created on first use, so importing the app does not require an API key. + return OpenAI() def generate_embedding(text: str) -> list[float]: - response = client.embeddings.create( - model="text-embedding-3-small", - input=text, - ) + response = get_client().embeddings.create(model=EMBEDDING_MODEL, input=text) return response.data[0].embedding def dot_product(a: list[float], b: list[float]) -> float: - return sum(x * y for x, y in zip(a, b)) + return sum(x * y for x, y in zip(a, b, strict=True)) def build_chunk_embeddings(chunks: list[str]) -> list[dict]: - chunk_embeddings = [] - - for chunk in chunks: - emb = generate_embedding(chunk) - chunk_embeddings.append({ - "text": chunk, - "embedding": emb - }) - - return chunk_embeddings + return [{"text": chunk, "embedding": generate_embedding(chunk)} for chunk in chunks] -def rank_chunks_by_similarity(query_embedding: list[float], chunk_embeddings: list[dict]) -> list[dict]: - scores = [] - - for item in chunk_embeddings: - score = dot_product(query_embedding, item["embedding"]) - scores.append({ - "text": item["text"], - "score": score - }) - +def rank_chunks_by_similarity( + query_embedding: list[float], chunk_embeddings: list[dict] +) -> list[dict]: + scores = [ + {"text": item["text"], "score": dot_product(query_embedding, item["embedding"])} + for item in chunk_embeddings + ] scores.sort(key=lambda x: x["score"], reverse=True) return scores def load_or_create_chunk_embeddings( - knowledge_path: str, - cache_path: str, - chunking_function + knowledge_path: str, cache_path: str, chunking_function ) -> list[dict]: knowledge_file = Path(knowledge_path) cache_file = Path(cache_path) @@ -56,7 +47,6 @@ def load_or_create_chunk_embeddings( if cache_file.exists(): cache_data = json.loads(cache_file.read_text(encoding="utf-8")) - if cache_data["last_modified"] == current_mtime: return cache_data["chunk_embeddings"] @@ -68,12 +58,8 @@ def load_or_create_chunk_embeddings( cache_data = { "last_modified": current_mtime, "chunks": chunks, - "chunk_embeddings": chunk_embeddings + "chunk_embeddings": chunk_embeddings, } + cache_file.write_text(json.dumps(cache_data, ensure_ascii=False, indent=2), encoding="utf-8") - cache_file.write_text( - json.dumps(cache_data, ensure_ascii=False, indent=2), - encoding="utf-8" - ) - - return chunk_embeddings \ No newline at end of file + return chunk_embeddings diff --git a/app/knowledge.py b/app/knowledge.py index 9eae89f..0126ea2 100644 --- a/app/knowledge.py +++ b/app/knowledge.py @@ -1,30 +1,14 @@ from pathlib import Path -def load_knowledge_base() -> str: - base_path = Path(__file__).resolve().parent.parent - file_path = base_path / "data" / "knowledge_base.txt" - return file_path.read_text(encoding="utf-8") - - def load_system_prompt() -> str: base_path = Path(__file__).resolve().parent.parent file_path = base_path / "prompts" / "system_prompt.txt" return file_path.read_text(encoding="utf-8") -def chunk_text(text: str, chunk_size: int = 500) -> list[str]: - chunks = [] - start = 0 - - while start < len(text): - end = start + chunk_size - chunks.append(text[start:end]) - start = end - - return chunks - def chunk_by_sections(text: str) -> list[str]: + """Split Markdown text into chunks, starting a new chunk at each `# ` or `## ` heading.""" chunks = [] current_chunk = "" @@ -41,4 +25,4 @@ def chunk_by_sections(text: str) -> list[str]: if current_chunk: chunks.append(current_chunk.strip()) - return chunks \ No newline at end of file + return chunks diff --git a/app/llm.py b/app/llm.py index 38d4029..1d43588 100644 --- a/app/llm.py +++ b/app/llm.py @@ -1,41 +1,31 @@ +from functools import lru_cache + from openai import OpenAI -from typing import List, Dict -client = OpenAI() +CHAT_MODEL = "gpt-4o-mini" + +@lru_cache(maxsize=1) +def get_client() -> OpenAI: + # Created on first use, so importing the app does not require an API key. + return OpenAI() -def ask_llm(question: str, context: str, system_prompt: str, history: List[Dict[str, str]] = None) -> str: - if history is None: - history = [] - + +def ask_llm( + question: str, + context: str, + system_prompt: str, + history: list[dict[str, str]] | None = None, +) -> str: messages = [ - { - "role": "system", - "content": system_prompt, - }, - { - "role": "system", - "content": f"Contexto:\n{context}", - }, + {"role": "system", "content": system_prompt}, + {"role": "system", "content": f"Contexto:\n{context}"}, ] - - # Add history messages - for msg in history: - messages.append({ - "role": msg["role"], - "content": msg["content"] - }) - - # Add current question - messages.append({ - "role": "user", - "content": question, - }) - - response = client.chat.completions.create( - model="gpt-4o-mini", - messages=messages, - ) + for msg in history or []: + messages.append({"role": msg["role"], "content": msg["content"]}) + messages.append({"role": "user", "content": question}) + + response = get_client().chat.completions.create(model=CHAT_MODEL, messages=messages) answer = response.choices[0].message.content - return answer or "Não foi possível gerar uma resposta." \ No newline at end of file + return answer or "Não foi possível gerar uma resposta." diff --git a/app/main.py b/app/main.py index dc6f309..2929c2a 100644 --- a/app/main.py +++ b/app/main.py @@ -1,4 +1,4 @@ -# FastAPI app for the RAG chatbot. At import time it loads the knowledge base, +# FastAPI app for the RAG chatbot. On startup it loads the knowledge base, # splits it into chunks and embeds them (cached in data/chunk_embeddings.json). # # Endpoints: @@ -7,36 +7,37 @@ # LLM with the session history, stores both messages and returns the answer. # - GET /conversations/{session_id}: stored messages of a session. -# Import necessary libraries and modules -from pathlib import Path -from dotenv import load_dotenv import uuid - -load_dotenv() +from contextlib import asynccontextmanager +from pathlib import Path from fastapi import FastAPI -from app.models import QuestionRequest, Message -from app.knowledge import load_system_prompt, chunk_by_sections -from app.llm import ask_llm + from app.embedding import load_or_create_chunk_embeddings -from app.services.rag_service import get_context_from_question +from app.knowledge import chunk_by_sections, load_system_prompt +from app.llm import ask_llm +from app.schemas import QuestionRequest from app.services.database_service import db_service +from app.services.rag_service import get_context_from_question + +BASE_PATH = Path(__file__).resolve().parent.parent +KNOWLEDGE_PATH = str(BASE_PATH / "data" / "knowledge_base.txt") +CACHE_PATH = str(BASE_PATH / "data" / "chunk_embeddings.json") + + +@asynccontextmanager +async def lifespan(app: FastAPI): + app.state.chunk_embeddings = load_or_create_chunk_embeddings( + knowledge_path=KNOWLEDGE_PATH, + cache_path=CACHE_PATH, + chunking_function=chunk_by_sections, + ) + yield -# Initialize FastAPI app -app = FastAPI() -# Load knowledge base and create/load chunks and embeddings at startup -base_path = Path(__file__).resolve().parent.parent -knowledge_path = str(base_path / "data" / "knowledge_base.txt") -cache_path = str(base_path / "data" / "chunk_embeddings.json") +app = FastAPI(title="rag-chatbot", lifespan=lifespan) -chunk_embeddings = load_or_create_chunk_embeddings( - knowledge_path=knowledge_path, - cache_path=cache_path, - chunking_function=chunk_by_sections, -) -# Define API endpoints @app.get("/health") def health(): return {"status": "ok"} @@ -46,34 +47,23 @@ def health(): def get_conversation_history(session_id: str): """Retrieve conversation history for a given session""" history = db_service.get_conversation_history(session_id) - return {"session_id": session_id, "messages": [msg.dict() for msg in history]} + return {"session_id": session_id, "messages": [msg.model_dump() for msg in history]} @app.post("/ask") def ask(req: QuestionRequest): - # Generate session_id if not provided session_id = req.session_id or str(uuid.uuid4()) - # Load system prompt and get context from question system_prompt = load_system_prompt() - # Get conversation history from database if session_id exists - history = req.history or [] - if not history and session_id: - history = db_service.get_conversation_history(session_id) + # History from the request body, or from the database when absent or empty. + history = req.history or db_service.get_conversation_history(session_id) - # Get context from question using the RAG service - context, sources = get_context_from_question(req.question, chunk_embeddings) + context, sources = get_context_from_question(req.question, app.state.chunk_embeddings) - # Ask the LLM for an answer based on the question, context, system prompt, and history - answer = ask_llm(req.question, context, system_prompt, [msg.dict() for msg in history]) + answer = ask_llm(req.question, context, system_prompt, [msg.model_dump() for msg in history]) - # Save user message and assistant response to database db_service.save_message_to_conversation(session_id, "user", req.question) db_service.save_message_to_conversation(session_id, "assistant", answer) - return { - "answer": answer, - "sources": sources, - "session_id": session_id - } \ No newline at end of file + return {"answer": answer, "sources": sources, "session_id": session_id} diff --git a/app/models.py b/app/models.py index 5fab34d..1a24812 100644 --- a/app/models.py +++ b/app/models.py @@ -1,31 +1,23 @@ -from pydantic import BaseModel -from typing import List, Dict, Optional -from datetime import datetime -from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey -from sqlalchemy.orm import relationship -from sqlalchemy.ext.declarative import declarative_base +from datetime import UTC, datetime -Base = declarative_base() +from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, Text +from sqlalchemy.orm import declarative_base, relationship -class Message(BaseModel): - role: str # "user" or "assistant" - content: str +Base = declarative_base() -class QuestionRequest(BaseModel): - question: str - history: Optional[List[Message]] = None - session_id: Optional[str] = None +def _utcnow() -> datetime: + # Columns are DateTime without timezone; keep storing naive UTC as before. + return datetime.now(UTC).replace(tzinfo=None) -# Database models class Conversation(Base): __tablename__ = "conversations" id = Column(Integer, primary_key=True, index=True) session_id = Column(String, unique=True, index=True) - created_at = Column(DateTime, default=datetime.utcnow) - updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) + created_at = Column(DateTime, default=_utcnow) + updated_at = Column(DateTime, default=_utcnow, onupdate=_utcnow) messages = relationship("MessageDB", back_populates="conversation") @@ -37,6 +29,6 @@ class MessageDB(Base): conversation_id = Column(Integer, ForeignKey("conversations.id")) role = Column(String) content = Column(Text) - created_at = Column(DateTime, default=datetime.utcnow) + created_at = Column(DateTime, default=_utcnow) - conversation = relationship("Conversation", back_populates="messages") \ No newline at end of file + conversation = relationship("Conversation", back_populates="messages") diff --git a/app/schemas.py b/app/schemas.py new file mode 100644 index 0000000..fdf9591 --- /dev/null +++ b/app/schemas.py @@ -0,0 +1,12 @@ +from pydantic import BaseModel + + +class Message(BaseModel): + role: str # "user" or "assistant" + content: str + + +class QuestionRequest(BaseModel): + question: str + history: list[Message] | None = None + session_id: str | None = None diff --git a/app/services/database_service.py b/app/services/database_service.py index 2ce79e2..e4139da 100644 --- a/app/services/database_service.py +++ b/app/services/database_service.py @@ -1,78 +1,50 @@ -from sqlalchemy import create_engine -from sqlalchemy.orm import sessionmaker, Session -from typing import List, Optional import os -from app.models import Conversation, MessageDB, Message -class DatabaseService: - def __init__(self): - database_url = os.getenv("DATABASE_URL", "postgresql://user:password@localhost/rag_chatbot") - self.engine = create_engine(database_url) - self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine) +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker - def get_db(self) -> Session: - db = self.SessionLocal() - try: - return db - finally: - db.close() +from app.models import Conversation, MessageDB +from app.schemas import Message - def create_conversation(self, session_id: str) -> Conversation: - db = self.get_db() - try: - conversation = Conversation(session_id=session_id) - db.add(conversation) - db.commit() - db.refresh(conversation) - return conversation - finally: - db.close() +DEFAULT_DATABASE_URL = "postgresql://user:password@localhost/rag_chatbot" - def get_conversation(self, session_id: str) -> Optional[Conversation]: - db = self.get_db() - try: - return db.query(Conversation).filter(Conversation.session_id == session_id).first() - finally: - db.close() - def add_message(self, conversation_id: int, role: str, content: str) -> MessageDB: - db = self.get_db() - try: - message = MessageDB(conversation_id=conversation_id, role=role, content=content) - db.add(message) - db.commit() - db.refresh(message) - return message - finally: - db.close() +class DatabaseService: + def __init__(self): + database_url = os.getenv("DATABASE_URL", DEFAULT_DATABASE_URL) + self.engine = create_engine(database_url) + self.SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=self.engine) - def get_conversation_history(self, session_id: str) -> List[Message]: - db = self.get_db() - try: - conversation = db.query(Conversation).filter(Conversation.session_id == session_id).first() + def get_conversation_history(self, session_id: str) -> list[Message]: + with self.SessionLocal() as db: + conversation = ( + db.query(Conversation).filter(Conversation.session_id == session_id).first() + ) if not conversation: return [] - messages = db.query(MessageDB).filter(MessageDB.conversation_id == conversation.id).order_by(MessageDB.created_at).all() - return [Message(role=msg.role, content=msg.content) for msg in messages] - finally: - db.close() + rows = ( + db.query(MessageDB) + .filter(MessageDB.conversation_id == conversation.id) + .order_by(MessageDB.created_at) + .all() + ) + return [Message(role=row.role, content=row.content) for row in rows] def save_message_to_conversation(self, session_id: str, role: str, content: str) -> None: - db = self.get_db() - try: - conversation = db.query(Conversation).filter(Conversation.session_id == session_id).first() + with self.SessionLocal() as db: + conversation = ( + db.query(Conversation).filter(Conversation.session_id == session_id).first() + ) if not conversation: conversation = Conversation(session_id=session_id) db.add(conversation) db.commit() db.refresh(conversation) - message = MessageDB(conversation_id=conversation.id, role=role, content=content) - db.add(message) + db.add(MessageDB(conversation_id=conversation.id, role=role, content=content)) db.commit() - finally: - db.close() -# Global instance -db_service = DatabaseService() \ No newline at end of file + +# Global instance; the engine only connects on first use. +db_service = DatabaseService() diff --git a/app/services/rag_service.py b/app/services/rag_service.py index 29b4348..49d45e8 100644 --- a/app/services/rag_service.py +++ b/app/services/rag_service.py @@ -1,7 +1,9 @@ from app.embedding import generate_embedding, rank_chunks_by_similarity -def get_context_from_question(question: str, chunk_embeddings: list[dict]) -> tuple[str, list[dict]]: +def get_context_from_question( + question: str, chunk_embeddings: list[dict] +) -> tuple[str, list[dict]]: query_embedding = generate_embedding(question) scores = rank_chunks_by_similarity(query_embedding, chunk_embeddings) @@ -17,4 +19,4 @@ def get_context_from_question(question: str, chunk_embeddings: list[dict]) -> tu for item in top_items ] - return context, sources \ No newline at end of file + return context, sources diff --git a/init_db.py b/init_db.py deleted file mode 100644 index 13e4286..0000000 --- a/init_db.py +++ /dev/null @@ -1,37 +0,0 @@ -#!/usr/bin/env python3 -""" -Database initialization script for the RAG chatbot. -Creates the tables directly from the SQLAlchemy models (Base.metadata.create_all), -without recording an Alembic revision. Use `alembic upgrade head` for the versioned path. -""" - -import os -import sys -from pathlib import Path - -# Add the project root to Python path -project_root = Path(__file__).resolve().parent -sys.path.insert(0, str(project_root)) - -from dotenv import load_dotenv -load_dotenv() - -from app.models import Base -from sqlalchemy import create_engine - -def init_database(): - """Initialize the database by creating all tables.""" - database_url = os.getenv("DATABASE_URL", "postgresql://user:password@localhost/rag_chatbot") - - print(f"Initializing database with URL: {database_url}") - - # Create engine - engine = create_engine(database_url) - - # Create all tables - Base.metadata.create_all(bind=engine) - - print("Database tables created successfully!") - -if __name__ == "__main__": - init_database() \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..8dee806 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[tool.ruff] +line-length = 100 +target-version = "py311" +extend-exclude = ["alembic/versions"] + +[tool.ruff.lint] +select = ["E", "F", "I", "B", "UP", "W"] + +[tool.ruff.lint.per-file-ignores] +"alembic/env.py" = ["E402"] # Alembic's template imports after config setup + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["."] diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..633bdf9 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,4 @@ +-r requirements.txt +pytest +httpx +ruff diff --git a/requirements.txt b/requirements.txt index 63ec19f..54ff268 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,4 +5,4 @@ openai python-dotenv psycopg2-binary sqlalchemy -alembic \ No newline at end of file +alembic diff --git a/run.sh b/run.sh index b30ed89..37b6fa4 100755 --- a/run.sh +++ b/run.sh @@ -1,7 +1,7 @@ #!/bin/bash -# garante que está na raiz do projeto +# Run from the project root with the virtualenv in .venv cd "$(dirname "$0")" source .venv/bin/activate -uvicorn app.main:app --reload \ No newline at end of file +uvicorn app.main:app --reload diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 0000000..917b8d4 --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,63 @@ +"""Endpoint tests with the OpenAI client and the database replaced by fakes.""" + +import pytest +from fastapi.testclient import TestClient + +from app import main +from app.schemas import Message + +FAKE_CHUNKS = [ + {"text": "## Plano Basic\naté 5 usuários", "embedding": [1.0, 0.0]}, + {"text": "## Suporte\nhorário comercial", "embedding": [0.0, 1.0]}, +] + + +class FakeDatabase: + def __init__(self): + self.messages: dict[str, list[Message]] = {} + + def get_conversation_history(self, session_id): + return list(self.messages.get(session_id, [])) + + def save_message_to_conversation(self, session_id, role, content): + self.messages.setdefault(session_id, []).append(Message(role=role, content=content)) + + +@pytest.fixture +def client(monkeypatch): + fake_db = FakeDatabase() + monkeypatch.setattr(main, "load_or_create_chunk_embeddings", lambda **kwargs: FAKE_CHUNKS) + monkeypatch.setattr(main, "db_service", fake_db) + monkeypatch.setattr("app.services.rag_service.generate_embedding", lambda text: [1.0, 0.0]) + monkeypatch.setattr( + main, "ask_llm", lambda question, context, system_prompt, history: f"answer to {question}" + ) + with TestClient(main.app) as test_client: + test_client.fake_db = fake_db + yield test_client + + +def test_health(client): + assert client.get("/health").json() == {"status": "ok"} + + +def test_ask_returns_answer_sources_and_session(client): + response = client.post("/ask", json={"question": "Quantos usuários?"}) + assert response.status_code == 200 + body = response.json() + assert body["answer"] == "answer to Quantos usuários?" + assert body["sources"][0]["text_preview"].startswith("## Plano Basic") + assert body["sources"][0]["score"] == 1.0 + assert body["session_id"] + + +def test_ask_stores_both_messages_and_history_is_returned(client): + session_id = client.post("/ask", json={"question": "Q1"}).json()["session_id"] + history = client.get(f"/conversations/{session_id}").json() + assert history["session_id"] == session_id + assert [m["role"] for m in history["messages"]] == ["user", "assistant"] + assert history["messages"][0]["content"] == "Q1" + + +def test_ask_rejects_missing_question(client): + assert client.post("/ask", json={}).status_code == 422 diff --git a/tests/test_embedding.py b/tests/test_embedding.py new file mode 100644 index 0000000..50d7dff --- /dev/null +++ b/tests/test_embedding.py @@ -0,0 +1,62 @@ +import json + +import pytest + +from app.embedding import dot_product, load_or_create_chunk_embeddings, rank_chunks_by_similarity + + +def test_dot_product(): + assert dot_product([1.0, 2.0, 3.0], [4.0, 5.0, 6.0]) == 32.0 + + +def test_dot_product_rejects_length_mismatch(): + with pytest.raises(ValueError): + dot_product([1.0], [1.0, 2.0]) + + +def test_rank_orders_by_descending_score(): + chunks = [ + {"text": "low", "embedding": [0.0, 1.0]}, + {"text": "high", "embedding": [1.0, 0.0]}, + {"text": "mid", "embedding": [0.5, 0.5]}, + ] + ranked = rank_chunks_by_similarity([1.0, 0.0], chunks) + assert [r["text"] for r in ranked] == ["high", "mid", "low"] + assert ranked[0]["score"] == 1.0 + + +def test_cache_is_reused_when_mtime_unchanged(tmp_path, monkeypatch): + knowledge = tmp_path / "kb.txt" + knowledge.write_text("# A\ntext", encoding="utf-8") + cache = tmp_path / "cache.json" + cache.write_text( + json.dumps( + { + "last_modified": knowledge.stat().st_mtime, + "chunks": ["# A\ntext"], + "chunk_embeddings": [{"text": "# A\ntext", "embedding": [1.0]}], + } + ), + encoding="utf-8", + ) + + def fail(*args, **kwargs): + raise AssertionError("embedding API must not be called when the cache is valid") + + monkeypatch.setattr("app.embedding.generate_embedding", fail) + result = load_or_create_chunk_embeddings(str(knowledge), str(cache), lambda t: [t]) + assert result == [{"text": "# A\ntext", "embedding": [1.0]}] + + +def test_cache_is_rebuilt_when_missing(tmp_path, monkeypatch): + knowledge = tmp_path / "kb.txt" + knowledge.write_text("# A\ntext\n# B\nmore", encoding="utf-8") + cache = tmp_path / "cache.json" + + monkeypatch.setattr("app.embedding.generate_embedding", lambda text: [float(len(text))]) + from app.knowledge import chunk_by_sections + + result = load_or_create_chunk_embeddings(str(knowledge), str(cache), chunk_by_sections) + assert [r["text"] for r in result] == ["# A\ntext", "# B\nmore"] + assert cache.exists() + assert json.loads(cache.read_text(encoding="utf-8"))["chunks"] == ["# A\ntext", "# B\nmore"] diff --git a/tests/test_knowledge.py b/tests/test_knowledge.py new file mode 100644 index 0000000..c4cfc80 --- /dev/null +++ b/tests/test_knowledge.py @@ -0,0 +1,36 @@ +from pathlib import Path + +from app.knowledge import chunk_by_sections + +SAMPLE = """# Title + +## Section A +line a1 + +### Subsection A.1 +line a1.1 + +## Section B +line b1 +""" + + +def test_chunks_start_at_h1_and_h2_only(): + chunks = chunk_by_sections(SAMPLE) + assert [c.splitlines()[0] for c in chunks] == ["# Title", "## Section A", "## Section B"] + + +def test_h3_stays_inside_parent_chunk(): + chunks = chunk_by_sections(SAMPLE) + assert "### Subsection A.1" in chunks[1] + + +def test_empty_text_gives_no_chunks(): + assert chunk_by_sections("") == [] + + +def test_sample_knowledge_base_has_14_chunks(): + text = (Path(__file__).resolve().parent.parent / "data" / "knowledge_base.txt").read_text( + encoding="utf-8" + ) + assert len(chunk_by_sections(text)) == 14 From c2feb8cabc44ebbf7137df18376c4bd8b798c105 Mon Sep 17 00:00:00 2001 From: Matuck Date: Tue, 8 Sep 2026 13:10:28 -0300 Subject: [PATCH 4/4] CI: bump checkout and setup-python actions to current majors Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QVi7kKzGQhgPKX8sDzTu6i --- .github/workflows/tests.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 87c8c8e..20d01d7 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -11,8 +11,8 @@ jobs: matrix: python-version: ["3.11", "3.12", "3.13"] steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v5 + - uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - run: pip install -r requirements-dev.txt