Revert "feat(knowledgebase): replace litellm-pgvector connector with memory-retrieval"

This reverts commit abeadc49c8.

Restores vendor/litellm-pgvector/ and the vector_store_registry wiring
(in-band file_search tool-call support) at the user's request, after
re-confirming against docs.litellm.ai/docs/completion/knowledgebase and
litellm-pgvector's own README that pg_vector is still not an in-process
vector_store_registry backend -- it requires this same standalone
connector service either way, so there is no simpler 'native' path that
was missed. Trading back in: 793 lines of vendored code, the untested
Prisma migration, and the git-context build risk noted in VENDORED.md
(all flagged as unverified against real hardware in issue #24), in
exchange for the file_search in-band tool call memory-retrieval did not
support.

Conflicts resolved on top of later commits (Redis, update.sh key-minting
fold-in):
- .env.example / docs/memory-knowledgebase.md: kept the auto-mint-via-
  update.sh language, renamed MEMORY_RETRIEVAL_* back to
  LITELLM_PGVECTOR_*.
- scripts/generate-secrets.sh: left deleted -- its job was folded into
  update.sh in 24d749b, unrelated to this revert.
- scripts/update.sh: renamed the MEMORY_RETRIEVAL_* secret/mint calls to
  LITELLM_PGVECTOR_* to match.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018WHfjWrSEcGhCoeu6dQfDa
This commit is contained in:
2026-09-02 22:45:08 +02:00
co-authored by Claude-Bot
parent b4dc83949e
commit e7983f0710
21 changed files with 1348 additions and 153 deletions
+5 -5
View File
@@ -49,13 +49,13 @@ REDIS_PASSWORD=
UI_USERNAME=admin
UI_PASSWORD=
# --- Knowledgebase (pgvector + memory-retrieval, see docs/memory-knowledgebase.md) ---
# --- Knowledgebase (pgvector + litellm-pgvector, see docs/memory-knowledgebase.md) ---
# Random value, filled in automatically by ./scripts/update.sh — leave blank.
PGVECTOR_DB_PASSWORD=
# Auth key memory-retrieval requires on its own API (its SERVER_API_KEY).
# Auth key litellm-pgvector requires on its own API (its SERVER_API_KEY).
# Random value, filled in automatically by ./scripts/update.sh — leave blank.
MEMORY_RETRIEVAL_API_KEY=
# A virtual key memory-retrieval uses to call back into litellm for
LITELLM_PGVECTOR_API_KEY=
# A virtual key litellm-pgvector uses to call back into litellm for
# embeddings. Minted automatically by ./scripts/update.sh — leave blank.
# Manual fallback: docs/proxy-key-onboarding.md.
MEMORY_RETRIEVAL_EMBEDDING_KEY=
LITELLM_PGVECTOR_EMBEDDING_KEY=
-1
View File
@@ -4,4 +4,3 @@
# to be committed to this repo.
data/
.leankg/
__pycache__/
+25 -17
View File
@@ -238,30 +238,38 @@ services:
retries: 10
# LiteLLM's native knowledgebase/vector-store feature has no Qdrant backend
# (the qdrant service above only serves Open WebUI's own RAG/Memory), so
# this small in-repo service wraps langchain_postgres.PGVector directly
# against pgvector-db instead — simpler than a vendored third-party
# connector (see docs/research/langchain-pgvector-vs-litellm-pgvector.md,
# which replaced the earlier litellm-pgvector approach). It still calls
# back into litellm for embeddings, same gateway boundary as everything
# else in this stack.
memory-retrieval:
# (the qdrant service above only serves Open WebUI's own RAG/Memory) — this
# companion service (github.com/BerriAI/litellm-pgvector) is the only
# self-hosted path. No published image exists yet, so this builds from a
# vendored copy in vendor/litellm-pgvector/ (see that dir's README) rather
# than a remote git build context — the server's Docker/BuildKit couldn't
# do an authenticated-looking clone of a public github.com repo (fails
# with "could not read Username ... terminal prompts disabled"), and
# vendoring sidesteps needing that debugged. See
# docs/research/litellm-knowledgebase.md.
# ponytail: unverified against real hardware — Prisma migration behavior on
# first boot and the exact vector_store_registry field names for the
# pg_vector provider need a live smoke test. See issue #24.
litellm-pgvector:
build:
context: ./services/memory-retrieval
container_name: memory-retrieval
context: ./vendor/litellm-pgvector
container_name: litellm-pgvector
depends_on:
pgvector-db:
condition: service_healthy
litellm:
condition: service_healthy
environment:
- DATABASE_URL=postgresql+psycopg://litellm_pgvector:${PGVECTOR_DB_PASSWORD}@pgvector-db:5432/litellm_pgvector
- LITELLM_BASE_URL=http://litellm:4000/v1
# A virtual key for this workload — see docs/proxy-key-onboarding.md.
- LITELLM_API_KEY=${MEMORY_RETRIEVAL_EMBEDDING_KEY}
- EMBEDDING_MODEL=local-embedding
- SERVER_API_KEY=${MEMORY_RETRIEVAL_API_KEY}
- COLLECTION_NAME=memory-and-notes
- DATABASE_URL=postgresql://litellm_pgvector:${PGVECTOR_DB_PASSWORD}@pgvector-db:5432/litellm_pgvector
- SERVER_API_KEY=${LITELLM_PGVECTOR_API_KEY}
# Calls back into litellm for embeddings, same pattern as any other
# workload — see docs/proxy-key-onboarding.md for issuing this key.
- EMBEDDING__MODEL=local-embedding
- EMBEDDING__BASE_URL=http://litellm:4000
- EMBEDDING__API_KEY=${LITELLM_PGVECTOR_EMBEDDING_KEY}
- EMBEDDING__DIMENSIONS=768
expose:
- "8000"
restart: unless-stopped
networks: [ai-stack]
+14 -13
View File
@@ -2,7 +2,7 @@
Three gateway-level capabilities added on top of the [AI gateway/proxy](https://git.arthurerlich.de/haylan/LLM-Server/issues/9), so every client behind LiteLLM gets them — not just Open WebUI. See [issue #21](https://git.arthurerlich.de/haylan/LLM-Server/issues/21) for the rationale.
**Not yet verified on real hardware** — see [issue #24](https://git.arthurerlich.de/haylan/LLM-Server/issues/24).
**Not yet verified on real hardware** — see [issue #24](https://git.arthurerlich.de/haylan/LLM-Server/issues/24). In particular: `litellm-pgvector`'s Prisma migrations on first boot, and the exact `vector_store_registry` field names for the `pg_vector` provider.
## Web search (SearXNG)
@@ -21,23 +21,23 @@ Requires `SEARXNG_LAN_IP` set in `.env` so the `litellm` container can resolve `
## Knowledgebase (vector store / RAG)
LiteLLM's native knowledgebase feature has **no Qdrant backend** — the `qdrant` service in this stack only serves Open WebUI's own separate RAG/Memory feature and is unrelated to this — and no `langchain_postgres` backend either. Instead, a small in-repo service wraps LangChain's `PGVector` directly against its own Postgres+pgvector database (`pgvector-db`). This replaced an earlier attempt to vendor the third-party `litellm-pgvector` connector — see `docs/research/langchain-pgvector-vs-litellm-pgvector.md` for why. One consequence: the OpenAI-style `file_search` tool call on a `/chat/completions` request doesn't work here (no `vector_store_registry` entry) — query `memory-retrieval` directly instead (see below).
LiteLLM's native knowledgebase feature has **no Qdrant backend** — the `qdrant` service in this stack only serves Open WebUI's own separate RAG/Memory feature and is unrelated to this. The only self-hosted path is [litellm-pgvector](https://github.com/BerriAI/litellm-pgvector), a companion service backed by its own Postgres+pgvector database (`pgvector-db`), which this stack now runs alongside `litellm`. Full research: `docs/research/litellm-knowledgebase.md`.
New pieces:
- **`embedding-server`** — a second llama.cpp instance (small footprint, `nomic-embed-text-v1.5`) serving `/v1/embeddings`. The chat model isn't embedding-trained and llama.cpp serves one model per process, so this can't just be a flag on `llama-server`.
- **`pgvector-db`** — Postgres with the pgvector extension, separate from `litellm-db`.
- **`memory-retrieval`** (`services/memory-retrieval/`) — a ~90-line FastAPI app wrapping `langchain_postgres.PGVector`, exposing `POST /ingest` and `POST /query`. Calls back into `litellm` for embeddings, same gateway boundary as everything else here.
- `litellm-config.yaml`'s `local-embedding` model entry, which `memory-retrieval` calls through.
- **`litellm-pgvector`** — the connector service; no published image exists, so it's built from a vendored copy of the upstream repo at `vendor/litellm-pgvector/` (see that dir's `VENDORED.md`) — a remote git build context failed on the server's Docker/BuildKit setup.
- `litellm-config.yaml`'s `local-embedding` model entry and `vector_store_registry` block, tying it together.
### First-time setup
```bash
docker compose --profile tools run --rm downloader-embedding # fetch the embedding model
docker compose up -d embedding-server pgvector-db memory-retrieval
docker compose up -d embedding-server pgvector-db litellm-pgvector
```
`./scripts/update.sh` mints `MEMORY_RETRIEVAL_EMBEDDING_KEY` automatically (a `memory-retrieval` virtual key via LiteLLM's own API) if it's blank — it calls back into `litellm` for embeddings, same as any other workload. See `docs/proxy-key-onboarding.md` if a mint fails and it needs doing by hand.
`./scripts/update.sh` mints `LITELLM_PGVECTOR_EMBEDDING_KEY` automatically (a `litellm-pgvector` virtual key via LiteLLM's own API) if it's blank — it calls back into `litellm` for embeddings, same as any other workload. See `docs/proxy-key-onboarding.md` if a mint fails and it needs doing by hand.
### Loading memory into it
@@ -51,13 +51,14 @@ One chunk per fact/paragraph line, tagged with `source`/`section` metadata. Re-r
### Querying it
Directly against `memory-retrieval` — no in-band `file_search` tool call (see above):
Via the OpenAI Assistants-style `file_search` tool on a chat completion:
```bash
curl -X POST http://memory-retrieval:8000/query \
-H "Authorization: Bearer <MEMORY_RETRIEVAL_API_KEY>" \
-H "Content-Type: application/json" \
-d '{"query": "...", "k": 5}'
```json
{
"model": "qwen3.8-27b-local",
"messages": [...],
"tools": [{"type": "file_search", "vector_store_ids": ["memory-and-notes"]}]
}
```
A caller wanting RAG-augmented chat does the "search then send" pattern: query `/query`, splice the results into the prompt, then call `litellm` as normal.
or directly: `POST /v1/vector_stores/memory-and-notes/search` with `{"query": "..."}`.
+16 -8
View File
@@ -25,7 +25,7 @@ model_list:
- model_name: local-embedding
litellm_params:
# Served by the dedicated embedding-server (nomic-embed-text-v1.5), not
# the chat model — see docker-compose.yml. Called by memory-retrieval
# the chat model — see docker-compose.yml. Called by litellm-pgvector
# to embed knowledgebase content, and available directly at
# /v1/embeddings for anything else that wants it.
model: openai/local-embedding
@@ -44,13 +44,21 @@ search_tools:
search_provider: searxng
api_base: http://search.home/
# Knowledgebase / RAG lives outside LiteLLM's own registry now — see the
# memory-retrieval service (docker-compose.yml) and
# docs/research/langchain-pgvector-vs-litellm-pgvector.md. LiteLLM's native
# vector_store_registry has no Qdrant provider and no langchain_postgres
# provider either, so registering a store here isn't an option; callers
# query memory-retrieval's /query endpoint directly instead of an in-band
# file_search tool call.
# Knowledgebase / RAG, backed by the litellm-pgvector companion service (NOT
# Qdrant — LiteLLM's native vector-store feature has no Qdrant provider, see
# docs/research/litellm-knowledgebase.md). vector_store_id is this proxy's
# own identifier for the store, not assigned by a backend.
# ponytail: field names here (custom_llm_provider: pg_vector, api_base
# pointed at litellm-pgvector) are the best fit from the litellm-pgvector
# README, not confirmed against a running deploy yet — smoke-test before
# relying on it. See issue #24.
vector_store_registry:
- vector_store_name: memory-and-notes
litellm_params:
vector_store_id: "memory-and-notes"
custom_llm_provider: pg_vector
api_base: http://litellm-pgvector:8000
embedding_model: local-embedding
router_settings:
# ponytail: LiteLLM's request-prioritization scheduler is beta (see
+24 -13
View File
@@ -1,37 +1,48 @@
#!/usr/bin/env bash
# Loads data/memory.md and data/claude-legacy-memory.md into the
# memory-retrieval knowledgebase (POST /ingest).
# Loads data/memory.md and data/claude-legacy-memory.md into the LiteLLM
# knowledgebase (the "memory-and-notes" vector store, see litellm-config.yaml)
# via litellm-pgvector's batch-embeddings endpoint.
#
# ponytail: one chunk per non-empty, non-heading line — both source files are
# already one fact/paragraph per line (no hard-wrapping), so this needs no
# real chunking logic. Re-run after editing either file; there's no dedup,
# PGVector always inserts — clear the collection first if you need a clean
# real chunking logic. Re-run after editing either file; there's no dedup, so
# this appends duplicates on a second run against unchanged content — clear
# the store first (DELETE the vector_store_id's rows) if you need a clean
# reload.
set -euo pipefail
cd "$(dirname "$0")/.."
[ -f .env ] && set -a && . ./.env && set +a
: "${MEMORY_RETRIEVAL_API_KEY:?Set MEMORY_RETRIEVAL_API_KEY in .env first}"
MEMORY_RETRIEVAL_URL="${MEMORY_RETRIEVAL_URL:-http://localhost:8000}"
: "${LITELLM_PGVECTOR_API_KEY:?Set LITELLM_PGVECTOR_API_KEY in .env first}"
LITELLM_PGVECTOR_URL="${LITELLM_PGVECTOR_URL:-http://localhost:8000}"
VECTOR_STORE_ID="memory-and-notes"
# Must match litellm-config.yaml's vector_store_registry entry — the
# registry just points at a store the backend must already know about.
# Ignores failure if it already exists (no documented idempotency check).
curl -sf -X POST "${LITELLM_PGVECTOR_URL}/v1/vector_stores" \
-H "Authorization: Bearer ${LITELLM_PGVECTOR_API_KEY}" \
-H "Content-Type: application/json" \
-d "{\"name\": \"${VECTOR_STORE_ID}\"}" > /dev/null 2>&1 || true
ingest_file() {
local file="$1" section=""
local chunks="[]"
local batch="[]"
while IFS= read -r line; do
case "$line" in
"#"*) section="${line#\# }"; section="${section#\#\# }"; continue ;;
""|"---") continue ;;
esac
chunks=$(jq --arg content "$line" --arg source "$file" --arg section "$section" \
'. += [{"content": $content, "metadata": {"source": $source, "section": $section}}]' <<<"$chunks")
batch=$(jq --arg content "$line" --arg source "$file" --arg section "$section" \
'. += [{"content": $content, "metadata": {"source": $source, "section": $section}}]' <<<"$batch")
done < "$file"
echo "Ingesting $(jq 'length' <<<"$chunks") chunks from $file..."
curl -sf -X POST "${MEMORY_RETRIEVAL_URL}/ingest" \
-H "Authorization: Bearer ${MEMORY_RETRIEVAL_API_KEY}" \
echo "Ingesting $(jq 'length' <<<"$batch") chunks from $file..."
curl -sf -X POST "${LITELLM_PGVECTOR_URL}/v1/vector_stores/${VECTOR_STORE_ID}/embeddings/batch" \
-H "Authorization: Bearer ${LITELLM_PGVECTOR_API_KEY}" \
-H "Content-Type: application/json" \
-d "{\"chunks\": ${chunks}}" > /dev/null
-d "$batch" > /dev/null
}
ingest_file data/memory.md
+3 -3
View File
@@ -39,7 +39,7 @@ set_if_blank LITELLM_DB_PASSWORD "$(openssl rand -hex 32)"
set_if_blank REDIS_PASSWORD "$(openssl rand -hex 32)"
set_if_blank UI_PASSWORD "$(openssl rand -hex 16)"
set_if_blank PGVECTOR_DB_PASSWORD "$(openssl rand -hex 32)"
set_if_blank MEMORY_RETRIEVAL_API_KEY "$(openssl rand -hex 32)"
set_if_blank LITELLM_PGVECTOR_API_KEY "$(openssl rand -hex 32)"
echo "==> resolving SEARXNG_LAN_IP"
# search.home is a LAN mDNS/local-DNS name — resolvable from this host, just
@@ -66,7 +66,7 @@ docker compose build --pull
echo "==> bringing up litellm (needed to mint virtual keys below)"
docker compose up -d --wait litellm-db litellm
# OPENWEBUI_LITELLM_KEY / MEMORY_RETRIEVAL_EMBEDDING_KEY are per-workload
# OPENWEBUI_LITELLM_KEY / LITELLM_PGVECTOR_EMBEDDING_KEY are per-workload
# virtual keys, not random secrets — minted via LiteLLM's own API
# (docs/proxy-key-onboarding.md documents the manual Admin UI route; this is
# the same thing over the REST endpoint LITELLM_MASTER_KEY already
@@ -96,7 +96,7 @@ mint_key_if_blank() {
fi
}
mint_key_if_blank OPENWEBUI_LITELLM_KEY openwebui
mint_key_if_blank MEMORY_RETRIEVAL_EMBEDDING_KEY memory-retrieval
mint_key_if_blank LITELLM_PGVECTOR_EMBEDDING_KEY litellm-pgvector
set -a && . ./.env && set +a
echo "==> recreating changed services"
-8
View File
@@ -1,8 +0,0 @@
FROM python:3.11-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY main.py .
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
-79
View File
@@ -1,79 +0,0 @@
"""Thin FastAPI wrapper around langchain_postgres.PGVector, replacing the
vendored litellm-pgvector connector. See
docs/research/langchain-pgvector-vs-litellm-pgvector.md for the rationale
and docs/memory-knowledgebase.md for usage.
ponytail: one fixed collection (COLLECTION_NAME) — this stack only needs one
knowledgebase ("memory-and-notes"), not a multi-tenant store registry.
"""
import os
from fastapi import Depends, FastAPI, HTTPException
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from langchain_openai import OpenAIEmbeddings
from langchain_postgres import PGVector
from pydantic import BaseModel
DATABASE_URL = os.environ["DATABASE_URL"]
LITELLM_BASE_URL = os.environ["LITELLM_BASE_URL"]
LITELLM_API_KEY = os.environ["LITELLM_API_KEY"]
EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "local-embedding")
SERVER_API_KEY = os.environ["SERVER_API_KEY"]
COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "memory-and-notes")
app = FastAPI(title="memory-retrieval", version="1.0.0")
embeddings = OpenAIEmbeddings(
model=EMBEDDING_MODEL, base_url=LITELLM_BASE_URL, api_key=LITELLM_API_KEY
)
vector_store = PGVector(
embeddings=embeddings,
collection_name=COLLECTION_NAME,
connection=DATABASE_URL,
use_jsonb=True,
)
security = HTTPBearer()
def check_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
if credentials.credentials != SERVER_API_KEY:
raise HTTPException(status_code=401, detail="Invalid API key")
class Chunk(BaseModel):
content: str
metadata: dict = {}
class IngestRequest(BaseModel):
chunks: list[Chunk]
class QueryRequest(BaseModel):
query: str
k: int = 5
@app.get("/health")
def health():
return {"status": "ok"}
@app.post("/ingest", dependencies=[Depends(check_api_key)])
def ingest(req: IngestRequest):
texts = [c.content for c in req.chunks]
metadatas = [c.metadata for c in req.chunks]
ids = vector_store.add_texts(texts=texts, metadatas=metadatas)
return {"ingested": len(ids)}
@app.post("/query", dependencies=[Depends(check_api_key)])
def query(req: QueryRequest):
results = vector_store.similarity_search_with_relevance_scores(req.query, k=req.k)
return {
"results": [
{"content": doc.page_content, "metadata": doc.metadata, "score": score}
for doc, score in results
]
}
@@ -1,6 +0,0 @@
fastapi
uvicorn[standard]
langchain-postgres
langchain-openai
psycopg[binary]
pydantic
+4
View File
@@ -0,0 +1,4 @@
.env
__pycache__/*
venv/*
venv
+32
View File
@@ -0,0 +1,32 @@
FROM python:3.11-slim
# Set environment variables
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
ENV PYTHONPATH=/app
# Install system dependencies
RUN apt-get update && apt-get install -y \
build-essential \
curl \
postgresql-client \
&& rm -rf /var/lib/apt/lists/*
# Set work directory
WORKDIR /app
# Install Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# Copy project
COPY . .
# Generate Prisma client
RUN prisma generate
# Expose port
EXPOSE 8000
# Command to run the application
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2025 Berri AI
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.
+388
View File
@@ -0,0 +1,388 @@
# OpenAI Vector Stores API with PGVector
A FastAPI application that provides OpenAI-compatible vector store endpoints using PGVector and LiteLLM proxy for embeddings.
## Features
- 🔌 OpenAI-compatible API endpoints
- 🗄️ PGVector for efficient vector storage and similarity search
- 🎛️ Configurable database field mappings
- 🔄 LiteLLM proxy integration for any embedding model
- 🐳 Docker support
- ⚡ FastAPI with async support
## API Endpoints
### 1. Create Vector Store
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"name": "Support FAQ"
}'
```
### 2. List Vector Stores
```bash
# List all vector stores
curl -X GET \
http://localhost:8000/v1/vector_stores \
-H "Authorization: Bearer your-api-key"
# List with pagination (limit and after parameters)
curl -X GET \
"http://localhost:8000/v1/vector_stores?limit=10&after=vs_abc123" \
-H "Authorization: Bearer your-api-key"
```
### 3. Add Single Embedding to Vector Store
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores/vs_abc123/embeddings \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"content": "Our return policy allows returns within 30 days of purchase.",
"embedding": [0.1, 0.2, 0.3, ...],
"metadata": {
"category": "returns",
"source": "faq",
"id": "return_policy_1"
}
}'
```
### 4. Add Multiple Embeddings (Batch)
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores/vs_abc123/embeddings/batch \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"embeddings": [
{
"content": "Our return policy allows returns within 30 days of purchase.",
"embedding": [0.1, 0.2, 0.3, ...],
"metadata": {"category": "returns"}
},
{
"content": "Shipping is free for orders over $50.",
"embedding": [0.4, 0.5, 0.6, ...],
"metadata": {"category": "shipping"}
}
]
}'
```
### 5. Search Vector Store
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores/vs_abc123/search \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"query": "What is the return policy?",
"limit": 20,
"filters": {"category": "support"}
}'
```
## Configuration
### Environment Variables
Create a `.env` file with the following configuration:
```bash
# Database Configuration
DATABASE_URL="postgresql://username:password@localhost:5432/vectordb?schema=public"
# API Configuration
SERVER_API_KEY="your-api-key-here"
# Server Configuration
HOST="0.0.0.0"
PORT=8000
# LiteLLM Proxy Configuration
EMBEDDING__MODEL="text-embedding-ada-002"
EMBEDDING__BASE_URL="http://localhost:4000"
EMBEDDING__API_KEY="sk-1234"
EMBEDDING__DIMENSIONS=1536
# Database Field Configuration (optional)
DB_FIELDS__ID_FIELD="id"
DB_FIELDS__CONTENT_FIELD="content"
DB_FIELDS__METADATA_FIELD="metadata"
DB_FIELDS__EMBEDDING_FIELD="embedding"
DB_FIELDS__VECTOR_STORE_ID_FIELD="vector_store_id"
DB_FIELDS__CREATED_AT_FIELD="created_at"
```
### Database Field Mapping
You can customize the database field names by setting environment variables:
- `DB_FIELDS__ID_FIELD` - Primary key field (default: "id")
- `DB_FIELDS__CONTENT_FIELD` - Text content field (default: "content")
- `DB_FIELDS__METADATA_FIELD` - JSON metadata field (default: "metadata")
- `DB_FIELDS__EMBEDDING_FIELD` - Vector embedding field (default: "embedding")
- `DB_FIELDS__VECTOR_STORE_ID_FIELD` - Foreign key field (default: "vector_store_id")
- `DB_FIELDS__CREATED_AT_FIELD` - Timestamp field (default: "created_at")
### LiteLLM Proxy Configuration
The application uses LiteLLM proxy for embeddings. Configure it with:
- `EMBEDDING__MODEL` - Model name (e.g., "text-embedding-ada-002")
- `EMBEDDING__BASE_URL` - LiteLLM proxy URL (e.g., "http://localhost:4000")
- `EMBEDDING__API_KEY` - LiteLLM proxy API key
- `EMBEDDING__DIMENSIONS` - Embedding dimensions (default: 1536)
## Setup and Installation
### 1. Install Dependencies
```bash
pip install -r requirements.txt
```
### 2. Database Setup
```bash
# Generate Prisma client
prisma generate
# Run database migrations
prisma db push
```
### 3. Set up LiteLLM Proxy
Start LiteLLM proxy pointing to your preferred embedding model:
```bash
# Example: Start LiteLLM proxy for OpenAI
litellm --model text-embedding-ada-002 --port 4000
```
### 4. Run the Application
```bash
python main.py
```
Or using uvicorn directly:
```bash
uvicorn main:app --host 0.0.0.0 --port 8000 --reload
```
## Docker Deployment
### Build and run with Docker:
```bash
# Build the image
docker build -t vector-store-api .
# Run the container
docker run -p 8000:8000 --env-file .env vector-store-api
```
## Database Schema
The application uses two main tables:
### vector_stores
- `id` (string, primary key)
- `name` (string)
- `file_counts` (json)
- `status` (string)
- `usage_bytes` (integer)
- `created_at` (timestamp)
- `expires_after` (json, optional)
- `expires_at` (timestamp, optional)
- `last_active_at` (timestamp, optional)
- `metadata` (json, optional)
### embeddings
- `id` (string, primary key)
- `vector_store_id` (string, foreign key)
- `content` (string)
- `embedding` (vector(1536))
- `metadata` (json, optional)
- `created_at` (timestamp)
## Supported Models
Any embedding model supported by LiteLLM proxy can be used. Examples:
- OpenAI: `text-embedding-ada-002`, `text-embedding-3-small`, `text-embedding-3-large`
- Cohere: `embed-english-v3.0`, `embed-multilingual-v3.0`
- Voyage: `voyage-2`, `voyage-large-2`
- And many more...
## API Response Format
### Vector Store Response
```json
{
"id": "vs_abc123",
"object": "vector_store",
"created_at": 1699024800,
"name": "Support FAQ",
"usage_bytes": 0,
"file_counts": {
"in_progress": 0,
"completed": 0,
"failed": 0,
"cancelled": 0,
"total": 0
},
"status": "completed",
"metadata": {}
}
```
### Vector Store List Response
```json
{
"object": "list",
"data": [
{
"id": "vs_abc123",
"object": "vector_store",
"created_at": 1699024800,
"name": "Support FAQ",
"usage_bytes": 1024,
"file_counts": {"completed": 5, "total": 5},
"status": "completed",
"metadata": {}
}
],
"first_id": "vs_abc123",
"last_id": "vs_def456",
"has_more": false
}
```
### Search Response
```json
{
"object": "vector_store.search",
"data": [
{
"id": "emb_123",
"content": "Return policy text...",
"score": 0.95,
"metadata": {"category": "support"}
}
],
"usage": {
"total_tokens": 1
}
}
```
## Example Search Request
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores/vs_support_faq/search \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"query": "How do I return an item?",
"limit": 5,
"return_metadata": true
}'
```
## Health Check
```bash
curl http://localhost:8000/health
```
## Migrating Existing Data
If you have an existing database with embeddings and content, you can easily migrate using the embedding APIs:
### 1. Create Vector Store
First, create a vector store for your data:
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"name": "Migrated Data",
"metadata": {"source": "legacy_system"}
}'
```
### 2. Batch Insert Embeddings
Use the batch endpoint to efficiently insert multiple embeddings:
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores/vs_your_id/embeddings/batch \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"embeddings": [
{
"content": "Your text content here",
"embedding": [0.1, 0.2, 0.3, ...1536 dimensions...],
"metadata": {"source_id": "doc_123", "category": "support"}
}
]
}'
```
### 3. Migration Script Example
Here's a Python script example for migrating from an existing database:
```python
import psycopg2
import requests
import json
# Connect to your existing database
conn = psycopg2.connect("your_existing_db_url")
cur = conn.cursor()
# Fetch existing data
cur.execute("SELECT content, embedding, metadata FROM your_table")
rows = cur.fetchall()
# Prepare batch data
embeddings = []
for content, embedding, metadata in rows:
embeddings.append({
"content": content,
"embedding": embedding.tolist(), # Convert numpy array to list
"metadata": metadata or {}
})
# Send batch to API
response = requests.post(
"http://localhost:8000/v1/vector_stores/your_vector_store_id/embeddings/batch",
headers={
"Authorization": "Bearer your-api-key",
"Content-Type": "application/json"
},
json={"embeddings": embeddings}
)
print(f"Migrated {len(embeddings)} embeddings")
```
## License
MIT License
+8
View File
@@ -0,0 +1,8 @@
Vendored from https://github.com/BerriAI/litellm-pgvector at commit
`b553f84a32f580b4303297df5567f25912b59d93` (main, 2026-09-02) — no changes
made to the source. See `docker-compose.yml`'s `litellm-pgvector` service
comment for why this is vendored instead of built from a remote git context.
To update: `git clone https://github.com/BerriAI/litellm-pgvector.git`
somewhere, copy everything except `.git/` over this directory, update the
commit hash above, and run `./scripts/update.sh`.
+60
View File
@@ -0,0 +1,60 @@
from typing import Dict, Optional
from pydantic import BaseModel
from pydantic_settings import BaseSettings
class DatabaseFieldConfig(BaseModel):
"""Configuration for database field mappings"""
id_field: str = "id"
content_field: str = "content"
metadata_field: str = "metadata"
embedding_field: str = "embedding"
vector_store_id_field: str = "vector_store_id"
created_at_field: str = "created_at"
class EmbeddingConfig(BaseModel):
"""Configuration for embedding generation via LiteLLM proxy"""
model: str = "text-embedding-ada-002"
base_url: str = "http://localhost:4000" # LiteLLM proxy URL
api_key: str = "sk-1234" # LiteLLM proxy API key
dimensions: int = 1536
class Settings(BaseSettings):
"""Application settings"""
# Database configuration
database_url: str = "postgresql://username:password@localhost:5432/vectordb?schema=public"
# API configuration
server_api_key: str = "your-api-key-here"
port: int = 8000
host: str = "0.0.0.0"
# Database field mappings
db_fields: DatabaseFieldConfig = DatabaseFieldConfig()
# Embedding configuration
embedding: EmbeddingConfig = EmbeddingConfig()
class Config:
env_file = ".env"
env_nested_delimiter = "__"
case_sensitive = False
# Allow environment variables like:
# DB_FIELDS__ID_FIELD=custom_id
# EMBEDDING__MODEL=text-embedding-3-small
# EMBEDDING__API_BASE=https://api.openai.com/v1
@property
def table_names(self) -> Dict[str, str]:
"""Get table names"""
return {
"vector_stores": "vector_stores",
"embeddings": "embeddings"
}
# Global settings instance
settings = Settings()
+90
View File
@@ -0,0 +1,90 @@
from typing import List, Optional
from config import settings, EmbeddingConfig
from litellm.types.utils import EmbeddingResponse
import litellm
import logging
class EmbeddingService:
"""Service for generating embeddings using OpenAI SDK pointed at LiteLLM proxy"""
def __init__(self, config: Optional[EmbeddingConfig] = None):
self.config = config or settings.embedding
async def generate_embedding(self, text: str) -> List[float]:
"""
Generate embedding for a single text using LiteLLM proxy
Args:
text: Text to embed
Returns:
List of floats representing the embedding vector
"""
try:
response: EmbeddingResponse = await litellm.aembedding(
model=self.config.model,
input=[text],
api_base=self.config.base_url,
api_key=self.config.api_key
)
logging.debug(f"Embedding response: {response}")
# Extract embedding from response
embedding = response.data[0]["embedding"]
# Validate embedding dimensions
if len(embedding) != self.config.dimensions:
raise ValueError(
f"Expected embedding dimension {self.config.dimensions}, "
f"got {len(embedding)}"
)
return embedding
except Exception as e:
raise RuntimeError(f"Failed to generate embedding: {str(e)}")
async def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
"""
Generate embeddings for multiple texts
Args:
texts: List of texts to embed
Returns:
List of embedding vectors
"""
try:
# Generate embeddings using LiteLLM
response = await litellm.aembedding(
model=self.config.model,
input=texts,
api_base=self.config.base_url,
api_key=self.config.api_key
)
# Extract embeddings from response
embeddings = [item.embedding for item in response.data]
# Validate embedding dimensions
for i, embedding in enumerate(embeddings):
if len(embedding) != self.config.dimensions:
raise ValueError(
f"Expected embedding dimension {self.config.dimensions} for text {i}, "
f"got {len(embedding)}"
)
return embeddings
except Exception as e:
raise RuntimeError(f"Failed to generate embeddings: {str(e)}")
def update_config(self, new_config: EmbeddingConfig):
"""Update the embedding configuration"""
self.config = new_config
# Global embedding service instance
embedding_service = EmbeddingService()
+522
View File
@@ -0,0 +1,522 @@
import os
import asyncio
import time
from typing import List, Optional
from fastapi import FastAPI, HTTPException, Depends, Header
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.middleware.cors import CORSMiddleware
from prisma import Prisma
from dotenv import load_dotenv
from models import (
VectorStoreCreateRequest,
VectorStoreResponse,
VectorStoreSearchRequest,
VectorStoreSearchResponse,
SearchResult,
EmbeddingCreateRequest,
EmbeddingResponse,
EmbeddingBatchCreateRequest,
EmbeddingBatchCreateResponse,
VectorStoreListResponse,
ContentChunk
)
from config import settings
from embedding_service import embedding_service
load_dotenv()
app = FastAPI(
title="OpenAI Vector Stores API",
description="OpenAI-compatible Vector Stores API using PGVector",
version="1.0.0"
)
# CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Global Prisma client
db = Prisma()
security = HTTPBearer()
async def get_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
"""Validate API key from Authorization header"""
expected_key = settings.server_api_key
if credentials.credentials != expected_key:
raise HTTPException(status_code=401, detail="Invalid API key")
return credentials.credentials
@app.on_event("startup")
async def startup():
"""Connect to database on startup"""
await db.connect()
@app.on_event("shutdown")
async def shutdown():
"""Disconnect from database on shutdown"""
await db.disconnect()
async def generate_query_embedding(query: str) -> List[float]:
"""
Generate an embedding for the query using LiteLLM
"""
return await embedding_service.generate_embedding(query)
@app.post("/v1/vector_stores", response_model=VectorStoreResponse)
async def create_vector_store(
request: VectorStoreCreateRequest,
api_key: str = Depends(get_api_key)
):
"""
Create a new vector store.
"""
try:
# Use raw SQL to insert the vector store with configurable table/field names
vector_store_table = settings.table_names["vector_stores"]
result = await db.query_raw(
f"""
INSERT INTO {vector_store_table} (id, name, file_counts, status, usage_bytes, expires_after, metadata, created_at)
VALUES (gen_random_uuid(), $1, $2, $3, $4, $5, $6, NOW())
RETURNING id, name, file_counts, status, usage_bytes, expires_after, expires_at, last_active_at, metadata,
EXTRACT(EPOCH FROM created_at)::bigint as created_at_timestamp
""",
request.name,
{"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
"completed",
0,
request.expires_after,
request.metadata or {}
)
if not result:
raise HTTPException(status_code=500, detail="Failed to create vector store")
vector_store = result[0]
# Convert to response format
created_at = int(vector_store["created_at_timestamp"])
expires_at = int(vector_store["expires_at"].timestamp()) if vector_store.get("expires_at") else None
last_active_at = int(vector_store["last_active_at"].timestamp()) if vector_store.get("last_active_at") else None
return VectorStoreResponse(
id=vector_store["id"],
created_at=created_at,
name=vector_store["name"],
usage_bytes=vector_store["usage_bytes"] or 0,
file_counts=vector_store["file_counts"] or {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
status=vector_store["status"],
expires_after=vector_store["expires_after"],
expires_at=expires_at,
last_active_at=last_active_at,
metadata=vector_store["metadata"]
)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to create vector store: {str(e)}")
@app.get("/v1/vector_stores", response_model=VectorStoreListResponse)
async def list_vector_stores(
limit: Optional[int] = 20,
after: Optional[str] = None,
before: Optional[str] = None,
api_key: str = Depends(get_api_key)
):
"""
List vector stores with optional pagination.
"""
try:
limit = min(limit or 20, 100) # Cap at 100 results
vector_store_table = settings.table_names["vector_stores"]
# Build base query
base_query = f"""
SELECT id, name, file_counts, status, usage_bytes, expires_after, expires_at, last_active_at, metadata,
EXTRACT(EPOCH FROM created_at)::bigint as created_at_timestamp
FROM {vector_store_table}
"""
# Add pagination conditions
conditions = []
params = []
param_count = 1
if after:
conditions.append(f"id > ${param_count}")
params.append(after)
param_count += 1
if before:
conditions.append(f"id < ${param_count}")
params.append(before)
param_count += 1
if conditions:
base_query += " WHERE " + " AND ".join(conditions)
# Add ordering and limit
final_query = base_query + f" ORDER BY created_at DESC LIMIT {limit + 1}"
# Execute query
results = await db.query_raw(final_query, *params)
# Check if there are more results
has_more = len(results) > limit
if has_more:
results = results[:limit] # Remove extra result
# Convert to response format
vector_stores = []
for row in results:
created_at = int(row["created_at_timestamp"])
expires_at = int(row["expires_at"].timestamp()) if row.get("expires_at") else None
last_active_at = int(row["last_active_at"].timestamp()) if row.get("last_active_at") else None
vector_store = VectorStoreResponse(
id=row["id"],
created_at=created_at,
name=row["name"],
usage_bytes=row["usage_bytes"] or 0,
file_counts=row["file_counts"] or {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
status=row["status"],
expires_after=row["expires_after"],
expires_at=expires_at,
last_active_at=last_active_at,
metadata=row["metadata"]
)
vector_stores.append(vector_store)
# Determine first_id and last_id
first_id = vector_stores[0].id if vector_stores else None
last_id = vector_stores[-1].id if vector_stores else None
return VectorStoreListResponse(
data=vector_stores,
first_id=first_id,
last_id=last_id,
has_more=has_more
)
except Exception as e:
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=f"Failed to list vector stores: {str(e)}")
@app.post("/v1/vector_stores/{vector_store_id}/search", response_model=VectorStoreSearchResponse)
@app.post("/vector_stores/{vector_store_id}/search", response_model=VectorStoreSearchResponse)
async def search_vector_store(
vector_store_id: str,
request: VectorStoreSearchRequest,
api_key: str = Depends(get_api_key)
):
"""
Search a vector store for similar content.
"""
try:
# Check if vector store exists
vector_store_table = settings.table_names["vector_stores"]
vector_store_result = await db.query_raw(
f"SELECT id FROM {vector_store_table} WHERE id = $1",
vector_store_id
)
if not vector_store_result:
raise HTTPException(status_code=404, detail="Vector store not found")
# Generate embedding for query
query_embedding = await generate_query_embedding(request.query)
query_vector_str = "[" + ",".join(map(str, query_embedding)) + "]"
# Build the raw SQL query for vector similarity search
limit = min(request.limit or 20, 100) # Cap at 100 results
# Base query with vector similarity using cosine distance
# Use configurable field names
fields = settings.db_fields
table_name = settings.table_names["embeddings"]
# Build query with proper parameter placeholders for Prisma
param_count = 1
query_params = [query_vector_str, vector_store_id]
base_query = f"""
SELECT
{fields.id_field},
{fields.content_field},
{fields.metadata_field},
({fields.embedding_field} <=> ${param_count}::vector) as distance
FROM {table_name}
WHERE {fields.vector_store_id_field} = ${param_count + 1}
"""
param_count += 2
# Add metadata filters if provided
filter_conditions = []
if request.filters:
for key, value in request.filters.items():
filter_conditions.append(f"{fields.metadata_field}->>${param_count} = ${param_count + 1}")
query_params.extend([key, str(value)])
param_count += 2
if filter_conditions:
base_query += " AND " + " AND ".join(filter_conditions)
# Add ordering and limit
final_query = base_query + f" ORDER BY distance ASC LIMIT {limit}"
# Execute the query
results = await db.query_raw(final_query, *query_params)
# Convert results to SearchResult objects
search_results = []
for row in results:
# Convert distance to similarity score (1 - normalized_distance)
# Cosine distance ranges from 0 (identical) to 2 (opposite)
similarity_score = max(0, 1 - (row['distance'] / 2))
# Extract filename from metadata or use a default
metadata = row[fields.metadata_field] or {}
filename = metadata.get('filename', 'document.txt')
content_chunks = [ContentChunk(type="text", text=row[fields.content_field])]
result = SearchResult(
file_id=row[fields.id_field],
filename=filename,
score=similarity_score,
attributes=metadata if request.return_metadata else None,
content=content_chunks
)
search_results.append(result)
return VectorStoreSearchResponse(
search_query=request.query,
data=search_results,
has_more=False, # TODO: Implement pagination
next_page=None
)
except HTTPException:
raise
except Exception as e:
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")
@app.post("/v1/vector_stores/{vector_store_id}/embeddings", response_model=EmbeddingResponse)
async def create_embedding(
vector_store_id: str,
request: EmbeddingCreateRequest,
api_key: str = Depends(get_api_key)
):
"""
Add a single embedding to a vector store.
"""
try:
# Check if vector store exists
vector_store_table = settings.table_names["vector_stores"]
vector_store_result = await db.query_raw(
f"SELECT id FROM {vector_store_table} WHERE id = $1",
vector_store_id
)
if not vector_store_result:
raise HTTPException(status_code=404, detail="Vector store not found")
# Convert embedding to vector string format
embedding_vector_str = "[" + ",".join(map(str, request.embedding)) + "]"
# Insert embedding using configurable field names
fields = settings.db_fields
table_name = settings.table_names["embeddings"]
result = await db.query_raw(
f"""
INSERT INTO {table_name} ({fields.id_field}, {fields.vector_store_id_field}, {fields.content_field},
{fields.embedding_field}, {fields.metadata_field}, {fields.created_at_field})
VALUES (gen_random_uuid(), $1, $2, $3::vector, $4, NOW())
RETURNING {fields.id_field}, {fields.vector_store_id_field}, {fields.content_field},
{fields.metadata_field}, EXTRACT(EPOCH FROM {fields.created_at_field})::bigint as created_at_timestamp
""",
vector_store_id,
request.content,
embedding_vector_str,
request.metadata or {}
)
if not result:
raise HTTPException(status_code=500, detail="Failed to create embedding")
embedding = result[0]
# Update vector store statistics
await db.query_raw(
f"""
UPDATE {vector_store_table}
SET
file_counts = jsonb_set(
jsonb_set(
COALESCE(file_counts, '{{"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}}'::jsonb),
'{{completed}}',
(COALESCE(file_counts->>'completed', '0')::int + 1)::text::jsonb
),
'{{total}}',
(COALESCE(file_counts->>'total', '0')::int + 1)::text::jsonb
),
usage_bytes = COALESCE(usage_bytes, 0) + LENGTH($2),
last_active_at = NOW()
WHERE id = $1
""",
vector_store_id,
request.content
)
return EmbeddingResponse(
id=embedding[fields.id_field],
vector_store_id=embedding[fields.vector_store_id_field],
content=embedding[fields.content_field],
metadata=embedding[fields.metadata_field],
created_at=int(embedding["created_at_timestamp"])
)
except HTTPException:
raise
except Exception as e:
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=f"Failed to create embedding: {str(e)}")
@app.post("/v1/vector_stores/{vector_store_id}/embeddings/batch", response_model=EmbeddingBatchCreateResponse)
async def create_embeddings_batch(
vector_store_id: str,
request: EmbeddingBatchCreateRequest,
api_key: str = Depends(get_api_key)
):
"""
Add multiple embeddings to a vector store in batch.
"""
try:
# Check if vector store exists
vector_store_table = settings.table_names["vector_stores"]
vector_store_result = await db.query_raw(
f"SELECT id FROM {vector_store_table} WHERE id = $1",
vector_store_id
)
if not vector_store_result:
raise HTTPException(status_code=404, detail="Vector store not found")
if not request.embeddings:
raise HTTPException(status_code=400, detail="No embeddings provided")
# Prepare batch insert
fields = settings.db_fields
table_name = settings.table_names["embeddings"]
# Build VALUES clause for batch insert
values_clauses = []
params = []
param_count = 1
for embedding_req in request.embeddings:
embedding_vector_str = "[" + ",".join(map(str, embedding_req.embedding)) + "]"
values_clauses.append(f"(gen_random_uuid(), ${param_count}, ${param_count + 1}, ${param_count + 2}::vector, ${param_count + 3}, NOW())")
params.extend([
vector_store_id,
embedding_req.content,
embedding_vector_str,
embedding_req.metadata or {}
])
param_count += 4
values_clause = ", ".join(values_clauses)
# Execute batch insert
result = await db.query_raw(
f"""
INSERT INTO {table_name} ({fields.id_field}, {fields.vector_store_id_field}, {fields.content_field},
{fields.embedding_field}, {fields.metadata_field}, {fields.created_at_field})
VALUES {values_clause}
RETURNING {fields.id_field}, {fields.vector_store_id_field}, {fields.content_field},
{fields.metadata_field}, EXTRACT(EPOCH FROM {fields.created_at_field})::bigint as created_at_timestamp
""",
*params
)
if not result:
raise HTTPException(status_code=500, detail="Failed to create embeddings")
# Calculate total content length for usage bytes update
total_content_length = sum(len(emb.content) for emb in request.embeddings)
# Update vector store statistics
await db.query_raw(
f"""
UPDATE {vector_store_table}
SET
file_counts = jsonb_set(
jsonb_set(
COALESCE(file_counts, '{{"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}}'::jsonb),
'{{completed}}',
(COALESCE(file_counts->>'completed', '0')::int + $2)::text::jsonb
),
'{{total}}',
(COALESCE(file_counts->>'total', '0')::int + $2)::text::jsonb
),
usage_bytes = COALESCE(usage_bytes, 0) + $3,
last_active_at = NOW()
WHERE id = $1
""",
vector_store_id,
len(request.embeddings),
total_content_length
)
# Convert results to response format
embeddings = []
for row in result:
embeddings.append(EmbeddingResponse(
id=row[fields.id_field],
vector_store_id=row[fields.vector_store_id_field],
content=row[fields.content_field],
metadata=row[fields.metadata_field],
created_at=int(row["created_at_timestamp"])
))
return EmbeddingBatchCreateResponse(
data=embeddings,
created=int(time.time())
)
except HTTPException:
raise
except Exception as e:
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=f"Failed to create embeddings batch: {str(e)}")
@app.get("/health")
async def health_check():
"""Health check endpoint"""
return {"status": "healthy", "timestamp": int(time.time())}
if __name__ == "__main__":
import uvicorn
uvicorn.run("main:app", host=settings.host, port=settings.port, reload=True)
+86
View File
@@ -0,0 +1,86 @@
from typing import Optional, Dict, Any, List
from pydantic import BaseModel
from datetime import datetime
class VectorStoreCreateRequest(BaseModel):
name: str
file_ids: Optional[List[str]] = None
expires_after: Optional[Dict[str, Any]] = None
chunking_strategy: Optional[Dict[str, Any]] = None
metadata: Optional[Dict[str, Any]] = None
class VectorStoreResponse(BaseModel):
id: str
object: str = "vector_store"
created_at: int
name: str
usage_bytes: int
file_counts: Dict[str, int]
status: str
expires_after: Optional[Dict[str, Any]] = None
expires_at: Optional[int] = None
last_active_at: Optional[int] = None
metadata: Optional[Dict[str, Any]] = None
class VectorStoreSearchRequest(BaseModel):
query: str
limit: Optional[int] = 20
filters: Optional[Dict[str, Any]] = None
return_metadata: Optional[bool] = True
class ContentChunk(BaseModel):
type: str = "text"
text: str
class SearchResult(BaseModel):
file_id: str
filename: str
score: float
attributes: Optional[Dict[str, Any]] = None
content: List[ContentChunk]
class VectorStoreSearchResponse(BaseModel):
object: str = "vector_store.search_results.page"
search_query: str
data: List[SearchResult]
has_more: bool = False
next_page: Optional[str] = None
class EmbeddingCreateRequest(BaseModel):
content: str
embedding: List[float]
metadata: Optional[Dict[str, Any]] = None
class EmbeddingResponse(BaseModel):
id: str
object: str = "embedding"
vector_store_id: str
content: str
metadata: Optional[Dict[str, Any]] = None
created_at: int
class EmbeddingBatchCreateRequest(BaseModel):
embeddings: List[EmbeddingCreateRequest]
class EmbeddingBatchCreateResponse(BaseModel):
object: str = "embedding.batch"
data: List[EmbeddingResponse]
created: int
class VectorStoreListResponse(BaseModel):
object: str = "list"
data: List[VectorStoreResponse]
first_id: Optional[str] = None
last_id: Optional[str] = None
has_more: bool = False
+40
View File
@@ -0,0 +1,40 @@
// This is your Prisma schema file,
// learn more about it in the docs: https://pris.ly/d/prisma-schema
generator client {
provider = "prisma-client-py"
}
datasource db {
provider = "postgresql"
url = env("DATABASE_URL")
}
model VectorStore {
id String @id @default(cuid())
name String
file_counts Json?
status String @default("completed")
usage_bytes Int? @default(0)
created_at DateTime @default(now())
expires_after Json?
expires_at DateTime?
last_active_at DateTime?
metadata Json?
embeddings Embedding[]
@@map("vector_stores")
}
model Embedding {
id String @id @default(cuid())
vector_store_id String
content String
embedding Unsupported("vector(1536)")
metadata Json?
created_at DateTime @default(now())
vector_store VectorStore @relation(fields: [vector_store_id], references: [id], onDelete: Cascade)
@@map("embeddings")
}
+10
View File
@@ -0,0 +1,10 @@
fastapi==0.104.1
uvicorn[standard]==0.24.0
prisma==0.11.0
python-dotenv==1.0.0
pydantic>=2.5.0
psycopg2-binary==2.9.7
pgvector==0.2.4
python-multipart==0.0.6
litellm==1.74.3
pydantic-settings==2.1.0