haylanandClaude-Bot
abeadc49c8
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>
80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
"""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
|
|
]
|
|
}
|