feat(knowledgebase): replace litellm-pgvector connector with memory-retrieval
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 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
]
|
||||
}
|
||||
Reference in New Issue
Block a user