diff --git a/vendor/litellm-pgvector/main.py b/vendor/litellm-pgvector/main.py index 159bcd9..b1b52c1 100644 --- a/vendor/litellm-pgvector/main.py +++ b/vendor/litellm-pgvector/main.py @@ -85,14 +85,22 @@ async def create_vector_store( try: # Use raw SQL to insert the vector store with configurable table/field names vector_store_table = settings.table_names["vector_stores"] - + + # ponytail: honor a caller-supplied id (request.id) instead of + # always minting one — litellm's vector_store_registry addresses + # this store by a fixed id (see litellm-config.yaml), which never + # matched anything when this always generated a random UUID. + import uuid as _uuid + vector_store_id = request.id or str(_uuid.uuid4()) + 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, + VALUES ($1, $2, $3, $4, $5, $6, $7, 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 """, + vector_store_id, request.name, {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, "completed", diff --git a/vendor/litellm-pgvector/models.py b/vendor/litellm-pgvector/models.py index 70f4bff..ccd8667 100644 --- a/vendor/litellm-pgvector/models.py +++ b/vendor/litellm-pgvector/models.py @@ -5,6 +5,12 @@ from datetime import datetime class VectorStoreCreateRequest(BaseModel): name: str + # ponytail: not part of upstream litellm-pgvector — added locally so + # callers (litellm's vector_store_registry, scripts/ingest-memory.sh) + # can pin a human-readable id instead of getting a random UUID back. + # litellm-config.yaml's vector_store_registry addresses stores by a + # fixed vector_store_id, which only works if creation can honor it. + id: Optional[str] = None file_ids: Optional[List[str]] = None expires_after: Optional[Dict[str, Any]] = None chunking_strategy: Optional[Dict[str, Any]] = None