haylan/LLM-Server
1
0
Fork 0
Code Issues 9 Pull Requests Actions Packages Projects Releases Wiki Activity

fix(litellm-pgvector): vendor the source instead of a remote git build context

Browse Source 
The server's Docker/BuildKit couldn't do a git-context build of a public
github.com repo (fails with "could not read Username ... terminal prompts
disabled" — an auth-shaped error for what should be an anonymous clone).
Rather than debug that, vendor litellm-pgvector's small source tree
directly (vendor/litellm-pgvector/, see VENDORED.md for provenance/update
steps) and build from the local path.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
haylan
2026-09-02 21:50:45 +02:00
co-authored by Claude-Bot
parent b0e37b2f4c
commit 7247674a9f
13 changed files with 1270 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vendor/litellm-pgvector/.gitignore
+4
View File
@@ -0,0 +1,4 @@
.env
__pycache__/*
venv/*
venv
vendor/litellm-pgvector/Dockerfile
Vendored
+32
View File
@@ -0,0 +1,32 @@
FROM python:3.11-slim
# Set environment variables
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
ENV PYTHONPATH=/app
# Install system dependencies
RUN apt-get update && apt-get install -y \
build-essential \
curl \
postgresql-client \
&& rm -rf /var/lib/apt/lists/*
# Set work directory
WORKDIR /app
# Install Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# Copy project
COPY . .
# Generate Prisma client
RUN prisma generate
# Expose port
EXPOSE 8000
# Command to run the application
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
vendor/litellm-pgvector/LICENSE
Vendored
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2025 Berri AI
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
vendor/litellm-pgvector/README.md
Vendored
+388
View File
@@ -0,0 +1,388 @@
# OpenAI Vector Stores API with PGVector
A FastAPI application that provides OpenAI-compatible vector store endpoints using PGVector and LiteLLM proxy for embeddings.
## Features
- 🔌 OpenAI-compatible API endpoints
- 🗄️ PGVector for efficient vector storage and similarity search
- 🎛️ Configurable database field mappings
- 🔄 LiteLLM proxy integration for any embedding model
- 🐳 Docker support
- ⚡ FastAPI with async support
## API Endpoints
### 1. Create Vector Store
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"name": "Support FAQ"
}'
```
### 2. List Vector Stores
```bash
# List all vector stores
curl -X GET \
http://localhost:8000/v1/vector_stores \
-H "Authorization: Bearer your-api-key"
# List with pagination (limit and after parameters)
curl -X GET \
"http://localhost:8000/v1/vector_stores?limit=10&after=vs_abc123" \
-H "Authorization: Bearer your-api-key"
```
### 3. Add Single Embedding to Vector Store
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores/vs_abc123/embeddings \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"content": "Our return policy allows returns within 30 days of purchase.",
"embedding": [0.1, 0.2, 0.3, ...],
"metadata": {
"category": "returns",
"source": "faq",
"id": "return_policy_1"
}
}'
```
### 4. Add Multiple Embeddings (Batch)
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores/vs_abc123/embeddings/batch \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"embeddings": [
{
"content": "Our return policy allows returns within 30 days of purchase.",
"embedding": [0.1, 0.2, 0.3, ...],
"metadata": {"category": "returns"}
},
{
"content": "Shipping is free for orders over $50.",
"embedding": [0.4, 0.5, 0.6, ...],
"metadata": {"category": "shipping"}
}
]
}'
```
### 5. Search Vector Store
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores/vs_abc123/search \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"query": "What is the return policy?",
"limit": 20,
"filters": {"category": "support"}
}'
```
## Configuration
### Environment Variables
Create a `.env` file with the following configuration:
```bash
# Database Configuration
DATABASE_URL="postgresql://username:password@localhost:5432/vectordb?schema=public"
# API Configuration
SERVER_API_KEY="your-api-key-here"
# Server Configuration
HOST="0.0.0.0"
PORT=8000
# LiteLLM Proxy Configuration
EMBEDDING__MODEL="text-embedding-ada-002"
EMBEDDING__BASE_URL="http://localhost:4000"
EMBEDDING__API_KEY="sk-1234"
EMBEDDING__DIMENSIONS=1536
# Database Field Configuration (optional)
DB_FIELDS__ID_FIELD="id"
DB_FIELDS__CONTENT_FIELD="content"
DB_FIELDS__METADATA_FIELD="metadata"
DB_FIELDS__EMBEDDING_FIELD="embedding"
DB_FIELDS__VECTOR_STORE_ID_FIELD="vector_store_id"
DB_FIELDS__CREATED_AT_FIELD="created_at"
```
### Database Field Mapping
You can customize the database field names by setting environment variables:
- `DB_FIELDS__ID_FIELD` - Primary key field (default: "id")
- `DB_FIELDS__CONTENT_FIELD` - Text content field (default: "content")
- `DB_FIELDS__METADATA_FIELD` - JSON metadata field (default: "metadata")
- `DB_FIELDS__EMBEDDING_FIELD` - Vector embedding field (default: "embedding")
- `DB_FIELDS__VECTOR_STORE_ID_FIELD` - Foreign key field (default: "vector_store_id")
- `DB_FIELDS__CREATED_AT_FIELD` - Timestamp field (default: "created_at")
### LiteLLM Proxy Configuration
The application uses LiteLLM proxy for embeddings. Configure it with:
- `EMBEDDING__MODEL` - Model name (e.g., "text-embedding-ada-002")
- `EMBEDDING__BASE_URL` - LiteLLM proxy URL (e.g., "http://localhost:4000")
- `EMBEDDING__API_KEY` - LiteLLM proxy API key
- `EMBEDDING__DIMENSIONS` - Embedding dimensions (default: 1536)
## Setup and Installation
### 1. Install Dependencies
```bash
pip install -r requirements.txt
```
### 2. Database Setup
```bash
# Generate Prisma client
prisma generate
# Run database migrations
prisma db push
```
### 3. Set up LiteLLM Proxy
Start LiteLLM proxy pointing to your preferred embedding model:
```bash
# Example: Start LiteLLM proxy for OpenAI
litellm --model text-embedding-ada-002 --port 4000
```
### 4. Run the Application
```bash
python main.py
```
Or using uvicorn directly:
```bash
uvicorn main:app --host 0.0.0.0 --port 8000 --reload
```
## Docker Deployment
### Build and run with Docker:
```bash
# Build the image
docker build -t vector-store-api .
# Run the container
docker run -p 8000:8000 --env-file .env vector-store-api
```
## Database Schema
The application uses two main tables:
### vector_stores
- `id` (string, primary key)
- `name` (string)
- `file_counts` (json)
- `status` (string)
- `usage_bytes` (integer)
- `created_at` (timestamp)
- `expires_after` (json, optional)
- `expires_at` (timestamp, optional)
- `last_active_at` (timestamp, optional)
- `metadata` (json, optional)
### embeddings
- `id` (string, primary key)
- `vector_store_id` (string, foreign key)
- `content` (string)
- `embedding` (vector(1536))
- `metadata` (json, optional)
- `created_at` (timestamp)
## Supported Models
Any embedding model supported by LiteLLM proxy can be used. Examples:
- OpenAI: `text-embedding-ada-002`, `text-embedding-3-small`, `text-embedding-3-large`
- Cohere: `embed-english-v3.0`, `embed-multilingual-v3.0`
- Voyage: `voyage-2`, `voyage-large-2`
- And many more...
## API Response Format
### Vector Store Response
```json
{
"id": "vs_abc123",
"object": "vector_store",
"created_at": 1699024800,
"name": "Support FAQ",
"usage_bytes": 0,
"file_counts": {
"in_progress": 0,
"completed": 0,
"failed": 0,
"cancelled": 0,
"total": 0
},
"status": "completed",
"metadata": {}
}
```
### Vector Store List Response
```json
{
"object": "list",
"data": [
{
"id": "vs_abc123",
"object": "vector_store",
"created_at": 1699024800,
"name": "Support FAQ",
"usage_bytes": 1024,
"file_counts": {"completed": 5, "total": 5},
"status": "completed",
"metadata": {}
}
],
"first_id": "vs_abc123",
"last_id": "vs_def456",
"has_more": false
}
```
### Search Response
```json
{
"object": "vector_store.search",
"data": [
{
"id": "emb_123",
"content": "Return policy text...",
"score": 0.95,
"metadata": {"category": "support"}
}
],
"usage": {
"total_tokens": 1
}
}
```
## Example Search Request
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores/vs_support_faq/search \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"query": "How do I return an item?",
"limit": 5,
"return_metadata": true
}'
```
## Health Check
```bash
curl http://localhost:8000/health
```
## Migrating Existing Data
If you have an existing database with embeddings and content, you can easily migrate using the embedding APIs:
### 1. Create Vector Store
First, create a vector store for your data:
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"name": "Migrated Data",
"metadata": {"source": "legacy_system"}
}'
```
### 2. Batch Insert Embeddings
Use the batch endpoint to efficiently insert multiple embeddings:
```bash
curl -X POST \
http://localhost:8000/v1/vector_stores/vs_your_id/embeddings/batch \
-H "Authorization: Bearer your-api-key" \
-H "Content-Type: application/json" \
-d '{
"embeddings": [
{
"content": "Your text content here",
"embedding": [0.1, 0.2, 0.3, ...1536 dimensions...],
"metadata": {"source_id": "doc_123", "category": "support"}
}
]
}'
```
### 3. Migration Script Example
Here's a Python script example for migrating from an existing database:
```python
import psycopg2
import requests
import json
# Connect to your existing database
conn = psycopg2.connect("your_existing_db_url")
cur = conn.cursor()
# Fetch existing data
cur.execute("SELECT content, embedding, metadata FROM your_table")
rows = cur.fetchall()
# Prepare batch data
embeddings = []
for content, embedding, metadata in rows:
embeddings.append({
"content": content,
"embedding": embedding.tolist(), # Convert numpy array to list
"metadata": metadata or {}
})
# Send batch to API
response = requests.post(
"http://localhost:8000/v1/vector_stores/your_vector_store_id/embeddings/batch",
headers={
"Authorization": "Bearer your-api-key",
"Content-Type": "application/json"
},
json={"embeddings": embeddings}
)
print(f"Migrated {len(embeddings)} embeddings")
```
## License
MIT License
vendor/litellm-pgvector/VENDORED.md
Vendored
+8
View File
@@ -0,0 +1,8 @@
Vendored from https://github.com/BerriAI/litellm-pgvector at commit
`b553f84a32f580b4303297df5567f25912b59d93` (main, 2026-09-02) — no changes
made to the source. See `docker-compose.yml`'s `litellm-pgvector` service
comment for why this is vendored instead of built from a remote git context.
To update: `git clone https://github.com/BerriAI/litellm-pgvector.git`
somewhere, copy everything except `.git/` over this directory, update the
commit hash above, and run `./scripts/update.sh`.
vendor/litellm-pgvector/config.py
Vendored
+60
View File
@@ -0,0 +1,60 @@
from typing import Dict, Optional
from pydantic import BaseModel
from pydantic_settings import BaseSettings
class DatabaseFieldConfig(BaseModel):
"""Configuration for database field mappings"""
id_field: str = "id"
content_field: str = "content"
metadata_field: str = "metadata"
embedding_field: str = "embedding"
vector_store_id_field: str = "vector_store_id"
created_at_field: str = "created_at"
class EmbeddingConfig(BaseModel):
"""Configuration for embedding generation via LiteLLM proxy"""
model: str = "text-embedding-ada-002"
base_url: str = "http://localhost:4000" # LiteLLM proxy URL
api_key: str = "sk-1234" # LiteLLM proxy API key
dimensions: int = 1536
class Settings(BaseSettings):
"""Application settings"""
# Database configuration
database_url: str = "postgresql://username:password@localhost:5432/vectordb?schema=public"
# API configuration
server_api_key: str = "your-api-key-here"
port: int = 8000
host: str = "0.0.0.0"
# Database field mappings
db_fields: DatabaseFieldConfig = DatabaseFieldConfig()
# Embedding configuration
embedding: EmbeddingConfig = EmbeddingConfig()
class Config:
env_file = ".env"
env_nested_delimiter = "__"
case_sensitive = False
# Allow environment variables like:
# DB_FIELDS__ID_FIELD=custom_id
# EMBEDDING__MODEL=text-embedding-3-small
# EMBEDDING__API_BASE=https://api.openai.com/v1
@property
def table_names(self) -> Dict[str, str]:
"""Get table names"""
return {
"vector_stores": "vector_stores",
"embeddings": "embeddings"
}
# Global settings instance
settings = Settings()
vendor/litellm-pgvector/embedding_service.py
Vendored
+90
View File
@@ -0,0 +1,90 @@
from typing import List, Optional
from config import settings, EmbeddingConfig
from litellm.types.utils import EmbeddingResponse
import litellm
import logging
class EmbeddingService:
"""Service for generating embeddings using OpenAI SDK pointed at LiteLLM proxy"""
def __init__(self, config: Optional[EmbeddingConfig] = None):
self.config = config or settings.embedding
async def generate_embedding(self, text: str) -> List[float]:
"""
Generate embedding for a single text using LiteLLM proxy
Args:
text: Text to embed
Returns:
List of floats representing the embedding vector
"""
try:
response: EmbeddingResponse = await litellm.aembedding(
model=self.config.model,
input=[text],
api_base=self.config.base_url,
api_key=self.config.api_key
)
logging.debug(f"Embedding response: {response}")
# Extract embedding from response
embedding = response.data[0]["embedding"]
# Validate embedding dimensions
if len(embedding) != self.config.dimensions:
raise ValueError(
f"Expected embedding dimension {self.config.dimensions}, "
f"got {len(embedding)}"
)
return embedding
except Exception as e:
raise RuntimeError(f"Failed to generate embedding: {str(e)}")
async def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
"""
Generate embeddings for multiple texts
Args:
texts: List of texts to embed
Returns:
List of embedding vectors
"""
try:
# Generate embeddings using LiteLLM
response = await litellm.aembedding(
model=self.config.model,
input=texts,
api_base=self.config.base_url,
api_key=self.config.api_key
)
# Extract embeddings from response
embeddings = [item.embedding for item in response.data]
# Validate embedding dimensions
for i, embedding in enumerate(embeddings):
if len(embedding) != self.config.dimensions:
raise ValueError(
f"Expected embedding dimension {self.config.dimensions} for text {i}, "
f"got {len(embedding)}"
)
return embeddings
except Exception as e:
raise RuntimeError(f"Failed to generate embeddings: {str(e)}")
def update_config(self, new_config: EmbeddingConfig):
"""Update the embedding configuration"""
self.config = new_config
# Global embedding service instance
embedding_service = EmbeddingService()
vendor/litellm-pgvector/main.py
Vendored
+522
View File
@@ -0,0 +1,522 @@
import os
import asyncio
import time
from typing import List, Optional
from fastapi import FastAPI, HTTPException, Depends, Header
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.middleware.cors import CORSMiddleware
from prisma import Prisma
from dotenv import load_dotenv
from models import (
VectorStoreCreateRequest,
VectorStoreResponse,
VectorStoreSearchRequest,
VectorStoreSearchResponse,
SearchResult,
EmbeddingCreateRequest,
EmbeddingResponse,
EmbeddingBatchCreateRequest,
EmbeddingBatchCreateResponse,
VectorStoreListResponse,
ContentChunk
)
from config import settings
from embedding_service import embedding_service
load_dotenv()
app = FastAPI(
title="OpenAI Vector Stores API",
description="OpenAI-compatible Vector Stores API using PGVector",
version="1.0.0"
)
# CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Global Prisma client
db = Prisma()
security = HTTPBearer()
async def get_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
"""Validate API key from Authorization header"""
expected_key = settings.server_api_key
if credentials.credentials != expected_key:
raise HTTPException(status_code=401, detail="Invalid API key")
return credentials.credentials
@app.on_event("startup")
async def startup():
"""Connect to database on startup"""
await db.connect()
@app.on_event("shutdown")
async def shutdown():
"""Disconnect from database on shutdown"""
await db.disconnect()
async def generate_query_embedding(query: str) -> List[float]:
"""
Generate an embedding for the query using LiteLLM
"""
return await embedding_service.generate_embedding(query)
@app.post("/v1/vector_stores", response_model=VectorStoreResponse)
async def create_vector_store(
request: VectorStoreCreateRequest,
api_key: str = Depends(get_api_key)
):
"""
Create a new vector store.
"""
try:
# Use raw SQL to insert the vector store with configurable table/field names
vector_store_table = settings.table_names["vector_stores"]
result = await db.query_raw(
f"""
INSERT INTO {vector_store_table} (id, name, file_counts, status, usage_bytes, expires_after, metadata, created_at)
VALUES (gen_random_uuid(), $1, $2, $3, $4, $5, $6, NOW())
RETURNING id, name, file_counts, status, usage_bytes, expires_after, expires_at, last_active_at, metadata,
EXTRACT(EPOCH FROM created_at)::bigint as created_at_timestamp
""",
request.name,
{"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
"completed",
0,
request.expires_after,
request.metadata or {}
)
if not result:
raise HTTPException(status_code=500, detail="Failed to create vector store")
vector_store = result[0]
# Convert to response format
created_at = int(vector_store["created_at_timestamp"])
expires_at = int(vector_store["expires_at"].timestamp()) if vector_store.get("expires_at") else None
last_active_at = int(vector_store["last_active_at"].timestamp()) if vector_store.get("last_active_at") else None
return VectorStoreResponse(
id=vector_store["id"],
created_at=created_at,
name=vector_store["name"],
usage_bytes=vector_store["usage_bytes"] or 0,
file_counts=vector_store["file_counts"] or {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
status=vector_store["status"],
expires_after=vector_store["expires_after"],
expires_at=expires_at,
last_active_at=last_active_at,
metadata=vector_store["metadata"]
)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to create vector store: {str(e)}")
@app.get("/v1/vector_stores", response_model=VectorStoreListResponse)
async def list_vector_stores(
limit: Optional[int] = 20,
after: Optional[str] = None,
before: Optional[str] = None,
api_key: str = Depends(get_api_key)
):
"""
List vector stores with optional pagination.
"""
try:
limit = min(limit or 20, 100) # Cap at 100 results
vector_store_table = settings.table_names["vector_stores"]
# Build base query
base_query = f"""
SELECT id, name, file_counts, status, usage_bytes, expires_after, expires_at, last_active_at, metadata,
EXTRACT(EPOCH FROM created_at)::bigint as created_at_timestamp
FROM {vector_store_table}
"""
# Add pagination conditions
conditions = []
params = []
param_count = 1
if after:
conditions.append(f"id > ${param_count}")
params.append(after)
param_count += 1
if before:
conditions.append(f"id < ${param_count}")
params.append(before)
param_count += 1
if conditions:
base_query += " WHERE " + " AND ".join(conditions)
# Add ordering and limit
final_query = base_query + f" ORDER BY created_at DESC LIMIT {limit + 1}"
# Execute query
results = await db.query_raw(final_query, *params)
# Check if there are more results
has_more = len(results) > limit
if has_more:
results = results[:limit] # Remove extra result
# Convert to response format
vector_stores = []
for row in results:
created_at = int(row["created_at_timestamp"])
expires_at = int(row["expires_at"].timestamp()) if row.get("expires_at") else None
last_active_at = int(row["last_active_at"].timestamp()) if row.get("last_active_at") else None
vector_store = VectorStoreResponse(
id=row["id"],
created_at=created_at,
name=row["name"],
usage_bytes=row["usage_bytes"] or 0,
file_counts=row["file_counts"] or {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
status=row["status"],
expires_after=row["expires_after"],
expires_at=expires_at,
last_active_at=last_active_at,
metadata=row["metadata"]
)
vector_stores.append(vector_store)
# Determine first_id and last_id
first_id = vector_stores[0].id if vector_stores else None
last_id = vector_stores[-1].id if vector_stores else None
return VectorStoreListResponse(
data=vector_stores,
first_id=first_id,
last_id=last_id,
has_more=has_more
)
except Exception as e:
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=f"Failed to list vector stores: {str(e)}")
@app.post("/v1/vector_stores/{vector_store_id}/search", response_model=VectorStoreSearchResponse)
@app.post("/vector_stores/{vector_store_id}/search", response_model=VectorStoreSearchResponse)
async def search_vector_store(
vector_store_id: str,
request: VectorStoreSearchRequest,
api_key: str = Depends(get_api_key)
):
"""
Search a vector store for similar content.
"""
try:
# Check if vector store exists
vector_store_table = settings.table_names["vector_stores"]
vector_store_result = await db.query_raw(
f"SELECT id FROM {vector_store_table} WHERE id = $1",
vector_store_id
)
if not vector_store_result:
raise HTTPException(status_code=404, detail="Vector store not found")
# Generate embedding for query
query_embedding = await generate_query_embedding(request.query)
query_vector_str = "[" + ",".join(map(str, query_embedding)) + "]"
# Build the raw SQL query for vector similarity search
limit = min(request.limit or 20, 100) # Cap at 100 results
# Base query with vector similarity using cosine distance
# Use configurable field names
fields = settings.db_fields
table_name = settings.table_names["embeddings"]
# Build query with proper parameter placeholders for Prisma
param_count = 1
query_params = [query_vector_str, vector_store_id]
base_query = f"""
SELECT
{fields.id_field},
{fields.content_field},
{fields.metadata_field},
({fields.embedding_field} <=> ${param_count}::vector) as distance
FROM {table_name}
WHERE {fields.vector_store_id_field} = ${param_count + 1}
"""
param_count += 2
# Add metadata filters if provided
filter_conditions = []
if request.filters:
for key, value in request.filters.items():
filter_conditions.append(f"{fields.metadata_field}->>${param_count} = ${param_count + 1}")
query_params.extend([key, str(value)])
param_count += 2
if filter_conditions:
base_query += " AND " + " AND ".join(filter_conditions)
# Add ordering and limit
final_query = base_query + f" ORDER BY distance ASC LIMIT {limit}"
# Execute the query
results = await db.query_raw(final_query, *query_params)
# Convert results to SearchResult objects
search_results = []
for row in results:
# Convert distance to similarity score (1 - normalized_distance)
# Cosine distance ranges from 0 (identical) to 2 (opposite)
similarity_score = max(0, 1 - (row['distance'] / 2))
# Extract filename from metadata or use a default
metadata = row[fields.metadata_field] or {}
filename = metadata.get('filename', 'document.txt')
content_chunks = [ContentChunk(type="text", text=row[fields.content_field])]
result = SearchResult(
file_id=row[fields.id_field],
filename=filename,
score=similarity_score,
attributes=metadata if request.return_metadata else None,
content=content_chunks
)
search_results.append(result)
return VectorStoreSearchResponse(
search_query=request.query,
data=search_results,
has_more=False, # TODO: Implement pagination
next_page=None
)
except HTTPException:
raise
except Exception as e:
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")
@app.post("/v1/vector_stores/{vector_store_id}/embeddings", response_model=EmbeddingResponse)
async def create_embedding(
vector_store_id: str,
request: EmbeddingCreateRequest,
api_key: str = Depends(get_api_key)
):
"""
Add a single embedding to a vector store.
"""
try:
# Check if vector store exists
vector_store_table = settings.table_names["vector_stores"]
vector_store_result = await db.query_raw(
f"SELECT id FROM {vector_store_table} WHERE id = $1",
vector_store_id
)
if not vector_store_result:
raise HTTPException(status_code=404, detail="Vector store not found")
# Convert embedding to vector string format
embedding_vector_str = "[" + ",".join(map(str, request.embedding)) + "]"
# Insert embedding using configurable field names
fields = settings.db_fields
table_name = settings.table_names["embeddings"]
result = await db.query_raw(
f"""
INSERT INTO {table_name} ({fields.id_field}, {fields.vector_store_id_field}, {fields.content_field},
{fields.embedding_field}, {fields.metadata_field}, {fields.created_at_field})
VALUES (gen_random_uuid(), $1, $2, $3::vector, $4, NOW())
RETURNING {fields.id_field}, {fields.vector_store_id_field}, {fields.content_field},
{fields.metadata_field}, EXTRACT(EPOCH FROM {fields.created_at_field})::bigint as created_at_timestamp
""",
vector_store_id,
request.content,
embedding_vector_str,
request.metadata or {}
)
if not result:
raise HTTPException(status_code=500, detail="Failed to create embedding")
embedding = result[0]
# Update vector store statistics
await db.query_raw(
f"""
UPDATE {vector_store_table}
SET
file_counts = jsonb_set(
jsonb_set(
COALESCE(file_counts, '{{"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}}'::jsonb),
'{{completed}}',
(COALESCE(file_counts->>'completed', '0')::int + 1)::text::jsonb
),
'{{total}}',
(COALESCE(file_counts->>'total', '0')::int + 1)::text::jsonb
),
usage_bytes = COALESCE(usage_bytes, 0) + LENGTH($2),
last_active_at = NOW()
WHERE id = $1
""",
vector_store_id,
request.content
)
return EmbeddingResponse(
id=embedding[fields.id_field],
vector_store_id=embedding[fields.vector_store_id_field],
content=embedding[fields.content_field],
metadata=embedding[fields.metadata_field],
created_at=int(embedding["created_at_timestamp"])
)
except HTTPException:
raise
except Exception as e:
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=f"Failed to create embedding: {str(e)}")
@app.post("/v1/vector_stores/{vector_store_id}/embeddings/batch", response_model=EmbeddingBatchCreateResponse)
async def create_embeddings_batch(
vector_store_id: str,
request: EmbeddingBatchCreateRequest,
api_key: str = Depends(get_api_key)
):
"""
Add multiple embeddings to a vector store in batch.
"""
try:
# Check if vector store exists
vector_store_table = settings.table_names["vector_stores"]
vector_store_result = await db.query_raw(
f"SELECT id FROM {vector_store_table} WHERE id = $1",
vector_store_id
)
if not vector_store_result:
raise HTTPException(status_code=404, detail="Vector store not found")
if not request.embeddings:
raise HTTPException(status_code=400, detail="No embeddings provided")
# Prepare batch insert
fields = settings.db_fields
table_name = settings.table_names["embeddings"]
# Build VALUES clause for batch insert
values_clauses = []
params = []
param_count = 1
for embedding_req in request.embeddings:
embedding_vector_str = "[" + ",".join(map(str, embedding_req.embedding)) + "]"
values_clauses.append(f"(gen_random_uuid(), ${param_count}, ${param_count + 1}, ${param_count + 2}::vector, ${param_count + 3}, NOW())")
params.extend([
vector_store_id,
embedding_req.content,
embedding_vector_str,
embedding_req.metadata or {}
])
param_count += 4
values_clause = ", ".join(values_clauses)
# Execute batch insert
result = await db.query_raw(
f"""
INSERT INTO {table_name} ({fields.id_field}, {fields.vector_store_id_field}, {fields.content_field},
{fields.embedding_field}, {fields.metadata_field}, {fields.created_at_field})
VALUES {values_clause}
RETURNING {fields.id_field}, {fields.vector_store_id_field}, {fields.content_field},
{fields.metadata_field}, EXTRACT(EPOCH FROM {fields.created_at_field})::bigint as created_at_timestamp
""",
*params
)
if not result:
raise HTTPException(status_code=500, detail="Failed to create embeddings")
# Calculate total content length for usage bytes update
total_content_length = sum(len(emb.content) for emb in request.embeddings)
# Update vector store statistics
await db.query_raw(
f"""
UPDATE {vector_store_table}
SET
file_counts = jsonb_set(
jsonb_set(
COALESCE(file_counts, '{{"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}}'::jsonb),
'{{completed}}',
(COALESCE(file_counts->>'completed', '0')::int + $2)::text::jsonb
),
'{{total}}',
(COALESCE(file_counts->>'total', '0')::int + $2)::text::jsonb
),
usage_bytes = COALESCE(usage_bytes, 0) + $3,
last_active_at = NOW()
WHERE id = $1
""",
vector_store_id,
len(request.embeddings),
total_content_length
)
# Convert results to response format
embeddings = []
for row in result:
embeddings.append(EmbeddingResponse(
id=row[fields.id_field],
vector_store_id=row[fields.vector_store_id_field],
content=row[fields.content_field],
metadata=row[fields.metadata_field],
created_at=int(row["created_at_timestamp"])
))
return EmbeddingBatchCreateResponse(
data=embeddings,
created=int(time.time())
)
except HTTPException:
raise
except Exception as e:
import traceback
traceback.print_exc()
raise HTTPException(status_code=500, detail=f"Failed to create embeddings batch: {str(e)}")
@app.get("/health")
async def health_check():
"""Health check endpoint"""
return {"status": "healthy", "timestamp": int(time.time())}
if __name__ == "__main__":
import uvicorn
uvicorn.run("main:app", host=settings.host, port=settings.port, reload=True)
vendor/litellm-pgvector/models.py
Vendored
+86
View File
@@ -0,0 +1,86 @@
from typing import Optional, Dict, Any, List
from pydantic import BaseModel
from datetime import datetime
class VectorStoreCreateRequest(BaseModel):
name: str
file_ids: Optional[List[str]] = None
expires_after: Optional[Dict[str, Any]] = None
chunking_strategy: Optional[Dict[str, Any]] = None
metadata: Optional[Dict[str, Any]] = None
class VectorStoreResponse(BaseModel):
id: str
object: str = "vector_store"
created_at: int
name: str
usage_bytes: int
file_counts: Dict[str, int]
status: str
expires_after: Optional[Dict[str, Any]] = None
expires_at: Optional[int] = None
last_active_at: Optional[int] = None
metadata: Optional[Dict[str, Any]] = None
class VectorStoreSearchRequest(BaseModel):
query: str
limit: Optional[int] = 20
filters: Optional[Dict[str, Any]] = None
return_metadata: Optional[bool] = True
class ContentChunk(BaseModel):
type: str = "text"
text: str
class SearchResult(BaseModel):
file_id: str
filename: str
score: float
attributes: Optional[Dict[str, Any]] = None
content: List[ContentChunk]
class VectorStoreSearchResponse(BaseModel):
object: str = "vector_store.search_results.page"
search_query: str
data: List[SearchResult]
has_more: bool = False
next_page: Optional[str] = None
class EmbeddingCreateRequest(BaseModel):
content: str
embedding: List[float]
metadata: Optional[Dict[str, Any]] = None
class EmbeddingResponse(BaseModel):
id: str
object: str = "embedding"
vector_store_id: str
content: str
metadata: Optional[Dict[str, Any]] = None
created_at: int
class EmbeddingBatchCreateRequest(BaseModel):
embeddings: List[EmbeddingCreateRequest]
class EmbeddingBatchCreateResponse(BaseModel):
object: str = "embedding.batch"
data: List[EmbeddingResponse]
created: int
class VectorStoreListResponse(BaseModel):
object: str = "list"
data: List[VectorStoreResponse]
first_id: Optional[str] = None
last_id: Optional[str] = None
has_more: bool = False
vendor/litellm-pgvector/prisma/schema.prisma
Vendored
+40
View File
@@ -0,0 +1,40 @@
// This is your Prisma schema file,
// learn more about it in the docs: https://pris.ly/d/prisma-schema
generator client {
provider = "prisma-client-py"
}
datasource db {
provider = "postgresql"
url = env("DATABASE_URL")
}
model VectorStore {
id String @id @default(cuid())
name String
file_counts Json?
status String @default("completed")
usage_bytes Int? @default(0)
created_at DateTime @default(now())
expires_after Json?
expires_at DateTime?
last_active_at DateTime?
metadata Json?
embeddings Embedding[]
@@map("vector_stores")
}
model Embedding {
id String @id @default(cuid())
vector_store_id String
content String
embedding Unsupported("vector(1536)")
metadata Json?
created_at DateTime @default(now())
vector_store VectorStore @relation(fields: [vector_store_id], references: [id], onDelete: Cascade)
@@map("embeddings")
}
vendor/litellm-pgvector/requirements.txt
Vendored
+10
View File
@@ -0,0 +1,10 @@
fastapi==0.104.1
uvicorn[standard]==0.24.0
prisma==0.11.0
python-dotenv==1.0.0
pydantic>=2.5.0
psycopg2-binary==2.9.7
pgvector==0.2.4
python-multipart==0.0.6
litellm==1.74.3
pydantic-settings==2.1.0