From abeadc49c8fb54653006dff0be577fd89b7e6091 Mon Sep 17 00:00:00 2001 From: ArthurErlich Date: Wed, 2 Sep 2026 22:09:25 +0200 Subject: [PATCH] feat(knowledgebase): replace litellm-pgvector connector with memory-retrieval MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per docs/research/langchain-pgvector-vs-litellm-pgvector.md (issue #25): the vendored litellm-pgvector connector (793 lines, Prisma migrations, a fragile git-context build) is replaced by a ~90-line FastAPI service (services/memory-retrieval/) wrapping langchain_postgres.PGVector directly against pgvector-db. Same gateway boundary — it still calls litellm for embeddings, nothing talks to Postgres or the model directly except this service. - New services/memory-retrieval/ (main.py, Dockerfile, requirements.txt): POST /ingest, POST /query, GET /health. - docker-compose.yml: litellm-pgvector service replaced by memory-retrieval; pgvector-db and embedding-server untouched. - litellm-config.yaml: vector_store_registry block removed (no langchain_postgres provider exists to register against; callers query memory-retrieval directly instead of an in-band file_search tool call — that mechanism was never confirmed working per issue #24 anyway). - scripts/ingest-memory.sh rewritten for the new /ingest endpoint (same per-line chunking, no dedup). - .env vars renamed: LITELLM_PGVECTOR_API_KEY/LITELLM_PGVECTOR_EMBEDDING_KEY -> MEMORY_RETRIEVAL_API_KEY/MEMORY_RETRIEVAL_EMBEDDING_KEY. - vendor/litellm-pgvector/ removed entirely. - docs/memory-knowledgebase.md updated for the new setup/query flow. Co-Authored-By: Claude Sonnet 5 --- .env.example | 12 +- .gitignore | 1 + docker-compose.yml | 42 +- docs/memory-knowledgebase.md | 27 +- litellm-config.yaml | 24 +- scripts/generate-secrets.sh | 6 +- scripts/ingest-memory.sh | 37 +- services/memory-retrieval/Dockerfile | 8 + services/memory-retrieval/main.py | 79 +++ services/memory-retrieval/requirements.txt | 6 + vendor/litellm-pgvector/.gitignore | 4 - vendor/litellm-pgvector/Dockerfile | 32 -- vendor/litellm-pgvector/LICENSE | 21 - vendor/litellm-pgvector/README.md | 388 -------------- vendor/litellm-pgvector/VENDORED.md | 8 - vendor/litellm-pgvector/config.py | 60 --- vendor/litellm-pgvector/embedding_service.py | 90 ---- vendor/litellm-pgvector/main.py | 522 ------------------- vendor/litellm-pgvector/models.py | 86 --- vendor/litellm-pgvector/prisma/schema.prisma | 40 -- vendor/litellm-pgvector/requirements.txt | 10 - 21 files changed, 154 insertions(+), 1349 deletions(-) create mode 100644 services/memory-retrieval/Dockerfile create mode 100644 services/memory-retrieval/main.py create mode 100644 services/memory-retrieval/requirements.txt delete mode 100644 vendor/litellm-pgvector/.gitignore delete mode 100644 vendor/litellm-pgvector/Dockerfile delete mode 100644 vendor/litellm-pgvector/LICENSE delete mode 100644 vendor/litellm-pgvector/README.md delete mode 100644 vendor/litellm-pgvector/VENDORED.md delete mode 100644 vendor/litellm-pgvector/config.py delete mode 100644 vendor/litellm-pgvector/embedding_service.py delete mode 100644 vendor/litellm-pgvector/main.py delete mode 100644 vendor/litellm-pgvector/models.py delete mode 100644 vendor/litellm-pgvector/prisma/schema.prisma delete mode 100644 vendor/litellm-pgvector/requirements.txt diff --git a/.env.example b/.env.example index 09698bc..b8ada63 100644 --- a/.env.example +++ b/.env.example @@ -43,12 +43,12 @@ LITELLM_DB_PASSWORD= UI_USERNAME=admin UI_PASSWORD= -# --- Knowledgebase (pgvector, see docs/memory-knowledgebase.md) --- +# --- Knowledgebase (pgvector + memory-retrieval, see docs/memory-knowledgebase.md) --- # Required — random values, e.g. `openssl rand -hex 32`. PGVECTOR_DB_PASSWORD= -# Auth key litellm-pgvector requires on its own API (its SERVER_API_KEY). -LITELLM_PGVECTOR_API_KEY= -# A virtual key litellm-pgvector uses to call back into litellm for +# Auth key memory-retrieval requires on its own API (its SERVER_API_KEY). +MEMORY_RETRIEVAL_API_KEY= +# A virtual key memory-retrieval uses to call back into litellm for # embeddings — create it in the Admin UI like any other workload key -# (see docs/proxy-key-onboarding.md), name it "litellm-pgvector". -LITELLM_PGVECTOR_EMBEDDING_KEY= +# (see docs/proxy-key-onboarding.md), name it "memory-retrieval". +MEMORY_RETRIEVAL_EMBEDDING_KEY= diff --git a/.gitignore b/.gitignore index 74c31ad..75aa8c4 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ # to be committed to this repo. data/ .leankg/ +__pycache__/ diff --git a/docker-compose.yml b/docker-compose.yml index 91c0af7..528e4b3 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -214,38 +214,30 @@ 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) — 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: + # (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: build: - context: ./vendor/litellm-pgvector - container_name: litellm-pgvector + context: ./services/memory-retrieval + container_name: memory-retrieval depends_on: pgvector-db: condition: service_healthy litellm: condition: service_healthy environment: - - 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" + - 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 restart: unless-stopped networks: [ai-stack] diff --git a/docs/memory-knowledgebase.md b/docs/memory-knowledgebase.md index 8a4f2f7..c1cc90e 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). In particular: `litellm-pgvector`'s Prisma migrations on first boot, and the exact `vector_store_registry` field names for the `pg_vector` provider. +**Not yet verified on real hardware** — see [issue #24](https://git.arthurerlich.de/haylan/LLM-Server/issues/24). ## Web search (SearXNG) @@ -21,23 +21,23 @@ Requires `SEARXNG_LAN_IP` set in `.env` (SearXNG's stable LAN IP — use a stati ## 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. 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`. +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). 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`. -- **`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. +- **`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. ### First-time setup ```bash docker compose --profile tools run --rm downloader-embedding # fetch the embedding model -docker compose up -d embedding-server pgvector-db litellm-pgvector +docker compose up -d embedding-server pgvector-db memory-retrieval ``` -Create a `litellm-pgvector` virtual key in LiteLLM's Admin UI (per `docs/proxy-key-onboarding.md`) and set it as `LITELLM_PGVECTOR_EMBEDDING_KEY` in `.env` — the connector calls back into `litellm` for embeddings, same as any other workload. +Create a `memory-retrieval` virtual key in LiteLLM's Admin UI (per `docs/proxy-key-onboarding.md`) and set it as `MEMORY_RETRIEVAL_EMBEDDING_KEY` in `.env` — it calls back into `litellm` for embeddings, same as any other workload. ### Loading memory into it @@ -51,14 +51,13 @@ One chunk per fact/paragraph line, tagged with `source`/`section` metadata. Re-r ### Querying it -Via the OpenAI Assistants-style `file_search` tool on a chat completion: +Directly against `memory-retrieval` — no in-band `file_search` tool call (see above): -```json -{ - "model": "qwen3.8-27b-local", - "messages": [...], - "tools": [{"type": "file_search", "vector_store_ids": ["memory-and-notes"]}] -} +```bash +curl -X POST http://memory-retrieval:8000/query \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{"query": "...", "k": 5}' ``` -or directly: `POST /v1/vector_stores/memory-and-notes/search` with `{"query": "..."}`. +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. diff --git a/litellm-config.yaml b/litellm-config.yaml index 9639172..9ca6882 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 litellm-pgvector + # the chat model — see docker-compose.yml. Called by memory-retrieval # to embed knowledgebase content, and available directly at # /v1/embeddings for anything else that wants it. model: openai/local-embedding @@ -44,21 +44,13 @@ search_tools: search_provider: searxng api_base: http://search.home/ -# 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 +# 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. router_settings: # ponytail: LiteLLM's request-prioritization scheduler is beta (see diff --git a/scripts/generate-secrets.sh b/scripts/generate-secrets.sh index 2368c63..ce8339e 100755 --- a/scripts/generate-secrets.sh +++ b/scripts/generate-secrets.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash # Generates random values for the secrets docker-compose.yml requires # (LITELLM_MASTER_KEY, LITELLM_SALT_KEY, LITELLM_DB_PASSWORD, UI_PASSWORD, -# PGVECTOR_DB_PASSWORD, LITELLM_PGVECTOR_API_KEY) and writes them into .env — +# PGVECTOR_DB_PASSWORD, MEMORY_RETRIEVAL_API_KEY) and writes them into .env — # creating it from .env.example first if it doesn't exist. # # ponytail: only fills in blank values, never overwrites ones you've already @@ -27,6 +27,6 @@ set_if_blank LITELLM_SALT_KEY "$(openssl rand -hex 32)" set_if_blank LITELLM_DB_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 LITELLM_PGVECTOR_API_KEY "$(openssl rand -hex 32)" +set_if_blank MEMORY_RETRIEVAL_API_KEY "$(openssl rand -hex 32)" -echo "Done. Review .env, then set OPENWEBUI_LITELLM_KEY, LITELLM_PGVECTOR_EMBEDDING_KEY, and SEARXNG_LAN_IP per docs/proxy-key-onboarding.md and docs/memory-knowledgebase.md." +echo "Done. Review .env, then set OPENWEBUI_LITELLM_KEY, MEMORY_RETRIEVAL_EMBEDDING_KEY, and SEARXNG_LAN_IP per docs/proxy-key-onboarding.md and docs/memory-knowledgebase.md." diff --git a/scripts/ingest-memory.sh b/scripts/ingest-memory.sh index 5b62808..ee7940a 100644 --- a/scripts/ingest-memory.sh +++ b/scripts/ingest-memory.sh @@ -1,48 +1,37 @@ #!/usr/bin/env bash -# 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. +# Loads data/memory.md and data/claude-legacy-memory.md into the +# memory-retrieval knowledgebase (POST /ingest). # # 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, 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 +# 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 # reload. set -euo pipefail cd "$(dirname "$0")/.." [ -f .env ] && set -a && . ./.env && set +a -: "${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 +: "${MEMORY_RETRIEVAL_API_KEY:?Set MEMORY_RETRIEVAL_API_KEY in .env first}" +MEMORY_RETRIEVAL_URL="${MEMORY_RETRIEVAL_URL:-http://localhost:8000}" ingest_file() { local file="$1" section="" - local batch="[]" + local chunks="[]" while IFS= read -r line; do case "$line" in "#"*) section="${line#\# }"; section="${section#\#\# }"; continue ;; ""|"---") continue ;; esac - batch=$(jq --arg content "$line" --arg source "$file" --arg section "$section" \ - '. += [{"content": $content, "metadata": {"source": $source, "section": $section}}]' <<<"$batch") + chunks=$(jq --arg content "$line" --arg source "$file" --arg section "$section" \ + '. += [{"content": $content, "metadata": {"source": $source, "section": $section}}]' <<<"$chunks") done < "$file" - 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}" \ + echo "Ingesting $(jq 'length' <<<"$chunks") chunks from $file..." + curl -sf -X POST "${MEMORY_RETRIEVAL_URL}/ingest" \ + -H "Authorization: Bearer ${MEMORY_RETRIEVAL_API_KEY}" \ -H "Content-Type: application/json" \ - -d "$batch" > /dev/null + -d "{\"chunks\": ${chunks}}" > /dev/null } ingest_file data/memory.md diff --git a/services/memory-retrieval/Dockerfile b/services/memory-retrieval/Dockerfile new file mode 100644 index 0000000..0c0c1e6 --- /dev/null +++ b/services/memory-retrieval/Dockerfile @@ -0,0 +1,8 @@ +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 new file mode 100644 index 0000000..2fd0773 --- /dev/null +++ b/services/memory-retrieval/main.py @@ -0,0 +1,79 @@ +"""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 new file mode 100644 index 0000000..c06e6dc --- /dev/null +++ b/services/memory-retrieval/requirements.txt @@ -0,0 +1,6 @@ +fastapi +uvicorn[standard] +langchain-postgres +langchain-openai +psycopg[binary] +pydantic diff --git a/vendor/litellm-pgvector/.gitignore b/vendor/litellm-pgvector/.gitignore deleted file mode 100644 index 0bdf641..0000000 --- a/vendor/litellm-pgvector/.gitignore +++ /dev/null @@ -1,4 +0,0 @@ -.env -__pycache__/* -venv/* -venv \ No newline at end of file diff --git a/vendor/litellm-pgvector/Dockerfile b/vendor/litellm-pgvector/Dockerfile deleted file mode 100644 index e4fef14..0000000 --- a/vendor/litellm-pgvector/Dockerfile +++ /dev/null @@ -1,32 +0,0 @@ -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 deleted file mode 100644 index 8d5e2b6..0000000 --- a/vendor/litellm-pgvector/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -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 deleted file mode 100644 index 439dc58..0000000 --- a/vendor/litellm-pgvector/README.md +++ /dev/null @@ -1,388 +0,0 @@ -# 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 deleted file mode 100644 index fd7efac..0000000 --- a/vendor/litellm-pgvector/VENDORED.md +++ /dev/null @@ -1,8 +0,0 @@ -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 deleted file mode 100644 index e83e681..0000000 --- a/vendor/litellm-pgvector/config.py +++ /dev/null @@ -1,60 +0,0 @@ -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 deleted file mode 100644 index c9e4930..0000000 --- a/vendor/litellm-pgvector/embedding_service.py +++ /dev/null @@ -1,90 +0,0 @@ -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 deleted file mode 100644 index 159bcd9..0000000 --- a/vendor/litellm-pgvector/main.py +++ /dev/null @@ -1,522 +0,0 @@ -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 deleted file mode 100644 index 70f4bff..0000000 --- a/vendor/litellm-pgvector/models.py +++ /dev/null @@ -1,86 +0,0 @@ -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 deleted file mode 100644 index faafada..0000000 --- a/vendor/litellm-pgvector/prisma/schema.prisma +++ /dev/null @@ -1,40 +0,0 @@ -// 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 deleted file mode 100644 index 9ec8b63..0000000 --- a/vendor/litellm-pgvector/requirements.txt +++ /dev/null @@ -1,10 +0,0 @@ -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