diff --git a/.env.example b/.env.example index a226286..83db9b0 100644 --- a/.env.example +++ b/.env.example @@ -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= diff --git a/.gitignore b/.gitignore index 75aa8c4..74c31ad 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,3 @@ # to be committed to this repo. data/ .leankg/ -__pycache__/ diff --git a/docker-compose.yml b/docker-compose.yml index 7cf9bc0..53b1c75 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -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] diff --git a/docs/memory-knowledgebase.md b/docs/memory-knowledgebase.md index d4c08c3..5ad1668 100644 --- a/docs/memory-knowledgebase.md +++ b/docs/memory-knowledgebase.md @@ -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 " \ - -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": "..."}`. diff --git a/litellm-config.yaml b/litellm-config.yaml index 8b81690..fcbb3c2 100644 --- a/litellm-config.yaml +++ b/litellm-config.yaml @@ -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 diff --git a/scripts/ingest-memory.sh b/scripts/ingest-memory.sh index ee7940a..5b62808 100644 --- a/scripts/ingest-memory.sh +++ b/scripts/ingest-memory.sh @@ -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 diff --git a/scripts/update.sh b/scripts/update.sh index 2f13689..67755ea 100644 --- a/scripts/update.sh +++ b/scripts/update.sh @@ -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" diff --git a/services/memory-retrieval/Dockerfile b/services/memory-retrieval/Dockerfile deleted file mode 100644 index 0c0c1e6..0000000 --- a/services/memory-retrieval/Dockerfile +++ /dev/null @@ -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"] diff --git a/services/memory-retrieval/main.py b/services/memory-retrieval/main.py deleted file mode 100644 index 2fd0773..0000000 --- a/services/memory-retrieval/main.py +++ /dev/null @@ -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 - ] - } diff --git a/services/memory-retrieval/requirements.txt b/services/memory-retrieval/requirements.txt deleted file mode 100644 index c06e6dc..0000000 --- a/services/memory-retrieval/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ -fastapi -uvicorn[standard] -langchain-postgres -langchain-openai -psycopg[binary] -pydantic diff --git a/vendor/litellm-pgvector/.gitignore b/vendor/litellm-pgvector/.gitignore new file mode 100644 index 0000000..0bdf641 --- /dev/null +++ b/vendor/litellm-pgvector/.gitignore @@ -0,0 +1,4 @@ +.env +__pycache__/* +venv/* +venv \ No newline at end of file diff --git a/vendor/litellm-pgvector/Dockerfile b/vendor/litellm-pgvector/Dockerfile new file mode 100644 index 0000000..e4fef14 --- /dev/null +++ b/vendor/litellm-pgvector/Dockerfile @@ -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"] \ No newline at end of file diff --git a/vendor/litellm-pgvector/LICENSE b/vendor/litellm-pgvector/LICENSE new file mode 100644 index 0000000..8d5e2b6 --- /dev/null +++ b/vendor/litellm-pgvector/LICENSE @@ -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. diff --git a/vendor/litellm-pgvector/README.md b/vendor/litellm-pgvector/README.md new file mode 100644 index 0000000..439dc58 --- /dev/null +++ b/vendor/litellm-pgvector/README.md @@ -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 diff --git a/vendor/litellm-pgvector/VENDORED.md b/vendor/litellm-pgvector/VENDORED.md new file mode 100644 index 0000000..fd7efac --- /dev/null +++ b/vendor/litellm-pgvector/VENDORED.md @@ -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`. diff --git a/vendor/litellm-pgvector/config.py b/vendor/litellm-pgvector/config.py new file mode 100644 index 0000000..e83e681 --- /dev/null +++ b/vendor/litellm-pgvector/config.py @@ -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() \ No newline at end of file diff --git a/vendor/litellm-pgvector/embedding_service.py b/vendor/litellm-pgvector/embedding_service.py new file mode 100644 index 0000000..c9e4930 --- /dev/null +++ b/vendor/litellm-pgvector/embedding_service.py @@ -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() \ No newline at end of file diff --git a/vendor/litellm-pgvector/main.py b/vendor/litellm-pgvector/main.py new file mode 100644 index 0000000..159bcd9 --- /dev/null +++ b/vendor/litellm-pgvector/main.py @@ -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) \ No newline at end of file diff --git a/vendor/litellm-pgvector/models.py b/vendor/litellm-pgvector/models.py new file mode 100644 index 0000000..70f4bff --- /dev/null +++ b/vendor/litellm-pgvector/models.py @@ -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 \ No newline at end of file diff --git a/vendor/litellm-pgvector/prisma/schema.prisma b/vendor/litellm-pgvector/prisma/schema.prisma new file mode 100644 index 0000000..faafada --- /dev/null +++ b/vendor/litellm-pgvector/prisma/schema.prisma @@ -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") +} \ No newline at end of file diff --git a/vendor/litellm-pgvector/requirements.txt b/vendor/litellm-pgvector/requirements.txt new file mode 100644 index 0000000..9ec8b63 --- /dev/null +++ b/vendor/litellm-pgvector/requirements.txt @@ -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 \ No newline at end of file