Skip to content

RAG Example

This example demonstrates how to use PyTiDB to build a minimal RAG application.

  • Use Ollama to deploy local embedding model and LLM model
  • Use Streamlit to build a Web UI for the RAG application
  • Use PyTiDB to build a minimal RAG application

RAG application built with PyTiDB

RAG application built with PyTiDB

Prerequisites

How to run

Step 1: Prepare the inference API

Pull the embedding and LLM model via ollama CLI:

ollama pull mxbai-embed-large
ollama pull gemma3:4b
ollama run gemma3:4b

Test the /embed and /generate endpoints to make sure they are running:

curl http://localhost:11434/api/embed -d '{
  "model": "mxbai-embed-large",
  "input": "Llamas are members of the camelid family"
}'

curl http://localhost:11434/api/generate -d '{
  "model": "gemma3:4b",
  "prompt": "Hello, Who are you?"
}'

Step 2: Clone the repository to local

git clone https://github.com/pingcap/pytidb.git
cd pytidb/examples/rag/;

Step 3: Install the required packages and setup environment

python -m venv .venv
source .venv/bin/activate
pip install -r reqs.txt

Step 4: Set up environment to connect to database

Go to TiDB Cloud console and get the connection parameters, then set up the environment variable like this:

cat > .env <<EOF
TIDB_HOST={gateway-region}.prod.aws.tidbcloud.com
TIDB_PORT=4000
TIDB_USERNAME={prefix}.root
TIDB_PASSWORD={password}
TIDB_DATABASE=test
EOF

Step 5: Run the Streamlit app

streamlit run main.py

Step 6: Open the browser and visit http://localhost:8501

Troubleshooting

502 Bad Gateway Error

Try to disable the global proxy settings.

  • Source Code: View on GitHub

  • Category: Ai-Apps

  • Description: Build a RAG application that combines document retrieval with language generation.

🏠 Back to Demo Gallery