Revert "feat(knowledgebase): replace litellm-pgvector connector with memory-retrieval"
This reverts commitabeadc49c8. 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 in24d749b, 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:
+5
-5
@@ -49,13 +49,13 @@ REDIS_PASSWORD=
|
|||||||
UI_USERNAME=admin
|
UI_USERNAME=admin
|
||||||
UI_PASSWORD=
|
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.
|
# Random value, filled in automatically by ./scripts/update.sh — leave blank.
|
||||||
PGVECTOR_DB_PASSWORD=
|
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.
|
# Random value, filled in automatically by ./scripts/update.sh — leave blank.
|
||||||
MEMORY_RETRIEVAL_API_KEY=
|
LITELLM_PGVECTOR_API_KEY=
|
||||||
# A virtual key memory-retrieval uses to call back into litellm for
|
# A virtual key litellm-pgvector uses to call back into litellm for
|
||||||
# embeddings. Minted automatically by ./scripts/update.sh — leave blank.
|
# embeddings. Minted automatically by ./scripts/update.sh — leave blank.
|
||||||
# Manual fallback: docs/proxy-key-onboarding.md.
|
# Manual fallback: docs/proxy-key-onboarding.md.
|
||||||
MEMORY_RETRIEVAL_EMBEDDING_KEY=
|
LITELLM_PGVECTOR_EMBEDDING_KEY=
|
||||||
|
|||||||
@@ -4,4 +4,3 @@
|
|||||||
# to be committed to this repo.
|
# to be committed to this repo.
|
||||||
data/
|
data/
|
||||||
.leankg/
|
.leankg/
|
||||||
__pycache__/
|
|
||||||
|
|||||||
+25
-17
@@ -238,30 +238,38 @@ services:
|
|||||||
retries: 10
|
retries: 10
|
||||||
|
|
||||||
# LiteLLM's native knowledgebase/vector-store feature has no Qdrant backend
|
# LiteLLM's native knowledgebase/vector-store feature has no Qdrant backend
|
||||||
# (the qdrant service above only serves Open WebUI's own RAG/Memory), so
|
# (the qdrant service above only serves Open WebUI's own RAG/Memory) — this
|
||||||
# this small in-repo service wraps langchain_postgres.PGVector directly
|
# companion service (github.com/BerriAI/litellm-pgvector) is the only
|
||||||
# against pgvector-db instead — simpler than a vendored third-party
|
# self-hosted path. No published image exists yet, so this builds from a
|
||||||
# connector (see docs/research/langchain-pgvector-vs-litellm-pgvector.md,
|
# vendored copy in vendor/litellm-pgvector/ (see that dir's README) rather
|
||||||
# which replaced the earlier litellm-pgvector approach). It still calls
|
# than a remote git build context — the server's Docker/BuildKit couldn't
|
||||||
# back into litellm for embeddings, same gateway boundary as everything
|
# do an authenticated-looking clone of a public github.com repo (fails
|
||||||
# else in this stack.
|
# with "could not read Username ... terminal prompts disabled"), and
|
||||||
memory-retrieval:
|
# 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:
|
build:
|
||||||
context: ./services/memory-retrieval
|
context: ./vendor/litellm-pgvector
|
||||||
container_name: memory-retrieval
|
container_name: litellm-pgvector
|
||||||
depends_on:
|
depends_on:
|
||||||
pgvector-db:
|
pgvector-db:
|
||||||
condition: service_healthy
|
condition: service_healthy
|
||||||
litellm:
|
litellm:
|
||||||
condition: service_healthy
|
condition: service_healthy
|
||||||
environment:
|
environment:
|
||||||
- DATABASE_URL=postgresql+psycopg://litellm_pgvector:${PGVECTOR_DB_PASSWORD}@pgvector-db:5432/litellm_pgvector
|
- DATABASE_URL=postgresql://litellm_pgvector:${PGVECTOR_DB_PASSWORD}@pgvector-db:5432/litellm_pgvector
|
||||||
- LITELLM_BASE_URL=http://litellm:4000/v1
|
- SERVER_API_KEY=${LITELLM_PGVECTOR_API_KEY}
|
||||||
# A virtual key for this workload — see docs/proxy-key-onboarding.md.
|
# Calls back into litellm for embeddings, same pattern as any other
|
||||||
- LITELLM_API_KEY=${MEMORY_RETRIEVAL_EMBEDDING_KEY}
|
# workload — see docs/proxy-key-onboarding.md for issuing this key.
|
||||||
- EMBEDDING_MODEL=local-embedding
|
- EMBEDDING__MODEL=local-embedding
|
||||||
- SERVER_API_KEY=${MEMORY_RETRIEVAL_API_KEY}
|
- EMBEDDING__BASE_URL=http://litellm:4000
|
||||||
- COLLECTION_NAME=memory-and-notes
|
- EMBEDDING__API_KEY=${LITELLM_PGVECTOR_EMBEDDING_KEY}
|
||||||
|
- EMBEDDING__DIMENSIONS=768
|
||||||
|
expose:
|
||||||
|
- "8000"
|
||||||
restart: unless-stopped
|
restart: unless-stopped
|
||||||
networks: [ai-stack]
|
networks: [ai-stack]
|
||||||
|
|
||||||
|
|||||||
@@ -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.
|
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)
|
## Web search (SearXNG)
|
||||||
|
|
||||||
@@ -21,23 +21,23 @@ Requires `SEARXNG_LAN_IP` set in `.env` so the `litellm` container can resolve `
|
|||||||
|
|
||||||
## Knowledgebase (vector store / RAG)
|
## 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:
|
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`.
|
- **`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`.
|
- **`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-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, which `memory-retrieval` calls through.
|
- `litellm-config.yaml`'s `local-embedding` model entry and `vector_store_registry` block, tying it together.
|
||||||
|
|
||||||
### First-time setup
|
### First-time setup
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
docker compose --profile tools run --rm downloader-embedding # fetch the embedding model
|
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
|
### Loading memory into it
|
||||||
|
|
||||||
@@ -51,13 +51,14 @@ One chunk per fact/paragraph line, tagged with `source`/`section` metadata. Re-r
|
|||||||
|
|
||||||
### Querying it
|
### 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
|
```json
|
||||||
curl -X POST http://memory-retrieval:8000/query \
|
{
|
||||||
-H "Authorization: Bearer <MEMORY_RETRIEVAL_API_KEY>" \
|
"model": "qwen3.8-27b-local",
|
||||||
-H "Content-Type: application/json" \
|
"messages": [...],
|
||||||
-d '{"query": "...", "k": 5}'
|
"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
@@ -25,7 +25,7 @@ model_list:
|
|||||||
- model_name: local-embedding
|
- model_name: local-embedding
|
||||||
litellm_params:
|
litellm_params:
|
||||||
# Served by the dedicated embedding-server (nomic-embed-text-v1.5), not
|
# 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
|
# to embed knowledgebase content, and available directly at
|
||||||
# /v1/embeddings for anything else that wants it.
|
# /v1/embeddings for anything else that wants it.
|
||||||
model: openai/local-embedding
|
model: openai/local-embedding
|
||||||
@@ -44,13 +44,21 @@ search_tools:
|
|||||||
search_provider: searxng
|
search_provider: searxng
|
||||||
api_base: http://search.home/
|
api_base: http://search.home/
|
||||||
|
|
||||||
# Knowledgebase / RAG lives outside LiteLLM's own registry now — see the
|
# Knowledgebase / RAG, backed by the litellm-pgvector companion service (NOT
|
||||||
# memory-retrieval service (docker-compose.yml) and
|
# Qdrant — LiteLLM's native vector-store feature has no Qdrant provider, see
|
||||||
# docs/research/langchain-pgvector-vs-litellm-pgvector.md. LiteLLM's native
|
# docs/research/litellm-knowledgebase.md). vector_store_id is this proxy's
|
||||||
# vector_store_registry has no Qdrant provider and no langchain_postgres
|
# own identifier for the store, not assigned by a backend.
|
||||||
# provider either, so registering a store here isn't an option; callers
|
# ponytail: field names here (custom_llm_provider: pg_vector, api_base
|
||||||
# query memory-retrieval's /query endpoint directly instead of an in-band
|
# pointed at litellm-pgvector) are the best fit from the litellm-pgvector
|
||||||
# file_search tool call.
|
# 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:
|
router_settings:
|
||||||
# ponytail: LiteLLM's request-prioritization scheduler is beta (see
|
# ponytail: LiteLLM's request-prioritization scheduler is beta (see
|
||||||
|
|||||||
+24
-13
@@ -1,37 +1,48 @@
|
|||||||
#!/usr/bin/env bash
|
#!/usr/bin/env bash
|
||||||
# Loads data/memory.md and data/claude-legacy-memory.md into the
|
# Loads data/memory.md and data/claude-legacy-memory.md into the LiteLLM
|
||||||
# memory-retrieval knowledgebase (POST /ingest).
|
# 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
|
# 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
|
# 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,
|
# real chunking logic. Re-run after editing either file; there's no dedup, so
|
||||||
# PGVector always inserts — clear the collection first if you need a clean
|
# 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.
|
# reload.
|
||||||
set -euo pipefail
|
set -euo pipefail
|
||||||
cd "$(dirname "$0")/.."
|
cd "$(dirname "$0")/.."
|
||||||
|
|
||||||
[ -f .env ] && set -a && . ./.env && set +a
|
[ -f .env ] && set -a && . ./.env && set +a
|
||||||
|
|
||||||
: "${MEMORY_RETRIEVAL_API_KEY:?Set MEMORY_RETRIEVAL_API_KEY in .env first}"
|
: "${LITELLM_PGVECTOR_API_KEY:?Set LITELLM_PGVECTOR_API_KEY in .env first}"
|
||||||
MEMORY_RETRIEVAL_URL="${MEMORY_RETRIEVAL_URL:-http://localhost:8000}"
|
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() {
|
ingest_file() {
|
||||||
local file="$1" section=""
|
local file="$1" section=""
|
||||||
local chunks="[]"
|
local batch="[]"
|
||||||
while IFS= read -r line; do
|
while IFS= read -r line; do
|
||||||
case "$line" in
|
case "$line" in
|
||||||
"#"*) section="${line#\# }"; section="${section#\#\# }"; continue ;;
|
"#"*) section="${line#\# }"; section="${section#\#\# }"; continue ;;
|
||||||
""|"---") continue ;;
|
""|"---") continue ;;
|
||||||
esac
|
esac
|
||||||
chunks=$(jq --arg content "$line" --arg source "$file" --arg section "$section" \
|
batch=$(jq --arg content "$line" --arg source "$file" --arg section "$section" \
|
||||||
'. += [{"content": $content, "metadata": {"source": $source, "section": $section}}]' <<<"$chunks")
|
'. += [{"content": $content, "metadata": {"source": $source, "section": $section}}]' <<<"$batch")
|
||||||
done < "$file"
|
done < "$file"
|
||||||
|
|
||||||
echo "Ingesting $(jq 'length' <<<"$chunks") chunks from $file..."
|
echo "Ingesting $(jq 'length' <<<"$batch") chunks from $file..."
|
||||||
curl -sf -X POST "${MEMORY_RETRIEVAL_URL}/ingest" \
|
curl -sf -X POST "${LITELLM_PGVECTOR_URL}/v1/vector_stores/${VECTOR_STORE_ID}/embeddings/batch" \
|
||||||
-H "Authorization: Bearer ${MEMORY_RETRIEVAL_API_KEY}" \
|
-H "Authorization: Bearer ${LITELLM_PGVECTOR_API_KEY}" \
|
||||||
-H "Content-Type: application/json" \
|
-H "Content-Type: application/json" \
|
||||||
-d "{\"chunks\": ${chunks}}" > /dev/null
|
-d "$batch" > /dev/null
|
||||||
}
|
}
|
||||||
|
|
||||||
ingest_file data/memory.md
|
ingest_file data/memory.md
|
||||||
|
|||||||
+3
-3
@@ -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 REDIS_PASSWORD "$(openssl rand -hex 32)"
|
||||||
set_if_blank UI_PASSWORD "$(openssl rand -hex 16)"
|
set_if_blank UI_PASSWORD "$(openssl rand -hex 16)"
|
||||||
set_if_blank PGVECTOR_DB_PASSWORD "$(openssl rand -hex 32)"
|
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"
|
echo "==> resolving SEARXNG_LAN_IP"
|
||||||
# search.home is a LAN mDNS/local-DNS name — resolvable from this host, just
|
# 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)"
|
echo "==> bringing up litellm (needed to mint virtual keys below)"
|
||||||
docker compose up -d --wait litellm-db litellm
|
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
|
# virtual keys, not random secrets — minted via LiteLLM's own API
|
||||||
# (docs/proxy-key-onboarding.md documents the manual Admin UI route; this is
|
# (docs/proxy-key-onboarding.md documents the manual Admin UI route; this is
|
||||||
# the same thing over the REST endpoint LITELLM_MASTER_KEY already
|
# the same thing over the REST endpoint LITELLM_MASTER_KEY already
|
||||||
@@ -96,7 +96,7 @@ mint_key_if_blank() {
|
|||||||
fi
|
fi
|
||||||
}
|
}
|
||||||
mint_key_if_blank OPENWEBUI_LITELLM_KEY openwebui
|
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
|
set -a && . ./.env && set +a
|
||||||
|
|
||||||
echo "==> recreating changed services"
|
echo "==> recreating changed services"
|
||||||
|
|||||||
@@ -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"]
|
|
||||||
@@ -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
|
|
||||||
@@ -0,0 +1,4 @@
|
|||||||
|
.env
|
||||||
|
__pycache__/*
|
||||||
|
venv/*
|
||||||
|
venv
|
||||||
Vendored
+32
@@ -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"]
|
||||||
Vendored
+21
@@ -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.
|
||||||
Vendored
+388
@@ -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
|
||||||
Vendored
+8
@@ -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`.
|
||||||
Vendored
+60
@@ -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
@@ -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()
|
||||||
Vendored
+522
@@ -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)
|
||||||
Vendored
+86
@@ -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
@@ -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
@@ -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
|
||||||
Reference in New Issue
Block a user