Skip to content

Commit

Permalink
postgress pgvector
Browse files Browse the repository at this point in the history
Sanjaypranav committed Dec 24, 2024
1 parent 749699b commit 6ee33b2
Showing 8 changed files with 313 additions and 66 deletions.
8 changes: 7 additions & 1 deletion .env.example
Original file line number Diff line number Diff line change
@@ -2,4 +2,10 @@ OPENAI_API_KEY=<your_openai_api_key>
TAVILY_API_KEY=<your_tavily_api_key>
LANGFUSE_SECRET_KEY=<your_langfuse_secret_key>
LANGFUSE_PUBLIC_KEY=<your_langfuse_public_key>
LANGFUSE_HOST=<your_langfuse_host>
LANGFUSE_HOST=<your_langfuse_host>
POSTGRES_HOST=<your_postgres_host>
POSTGRES_PORT=<your_postgres_port>
POSTGRES_DB=<your_postgres_db>
POSTGRES_USER=<your_postgres_user>
PGADMIN_DEFAULT_EMAIL=prakash@sayvai.io
PGADMIN_DEFAULT_PASSWORD=sayvai
15 changes: 15 additions & 0 deletions Dockerfile.pgvector
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
# This is installing the pgvector extension for postgres
FROM postgres:latest

RUN apt-get update && apt-get install -y \
build-essential \
git \
postgresql-server-dev-all \
&& rm -rf /var/lib/apt/lists/*

WORKDIR /tmp
RUN git clone https://github.com/pgvector/pgvector.git

WORKDIR /tmp/pgvector
RUN make
RUN make install
57 changes: 56 additions & 1 deletion docker-compose.yml
Original file line number Diff line number Diff line change
@@ -11,5 +11,60 @@ services:
- LANGFUSE_HOST=${LANGFUSE_HOST}
- LANGFUSE_PUBLIC_KEY=${LANGFUSE_PUBLIC_KEY}
- LANGFUSE_SECRET_KEY=${LANGFUSE_SECRET_KEY}
# Database connection environment variables
- POSTGRES_HOST=${POSTGRES_HOST}
- POSTGRES_PORT=${POSTGRES_PORT}
- POSTGRES_DB=${POSTGRES_DB}
- POSTGRES_USER=${POSTGRES_USER}
- POSTGRES_PASSWORD=${POSTGRES_PASSWORD}
ports:
- "8000:8000"
- "8000:8000"
networks:
- app-network
depends_on:
postgres:
condition: service_healthy

postgres:
build:
context: .
dockerfile: Dockerfile.pgvector
environment:
- POSTGRES_HOST=${POSTGRES_HOST}
- POSTGRES_PORT=${POSTGRES_PORT}
- POSTGRES_DB=${POSTGRES_DB}
- POSTGRES_USER=${POSTGRES_USER}
- POSTGRES_PASSWORD=${POSTGRES_PASSWORD}
ports:
- "5432:5432"
volumes:
- postgres_data:/var/lib/postgresql/data
- ./init-scripts/:/docker-entrypoint-initdb.d/
networks:
- app-network
# command: [ "postgres", "-c", "shared_preload_libraries=vector" ]
healthcheck:
test: [ "CMD-SHELL", "pg_isready -U sayvai -d user-conversations" ]
interval: 10s
timeout: 5s
retries: 5
start_period: 30s

pgadmin:
image: dpage/pgadmin4:latest
environment:
- PGADMIN_DEFAULT_EMAIL=${PGADMIN_DEFAULT_EMAIL}
- PGADMIN_DEFAULT_PASSWORD=${PGADMIN_DEFAULT_PASSWORD}
ports:
- "5050:80"
depends_on:
- postgres
networks:
- app-network

networks:
app-network:
driver: bridge

volumes:
postgres_data:
8 changes: 8 additions & 0 deletions init-scripts/vector_extension.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
#!/bin/bash
set -e

echo "In create extension"

psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname="$POSTGRES_DB" <<EOFSQL
CREATE EXTENSION vector;
EOFSQL
255 changes: 196 additions & 59 deletions poetry.lock

Large diffs are not rendered by default.

2 changes: 2 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -17,6 +17,8 @@ fastapi = {extras = ["all"], version = "^0.115.6"}
langfuse = "^2.57.0"
langchain-groq = "^0.2.2"
langchain-ollama = "^0.2.2"
langchain-postgres = "^0.0.12"
psycopg2-binary = "^2.9.10"


[build-system]
2 changes: 2 additions & 0 deletions travis/agent.py
Original file line number Diff line number Diff line change
@@ -4,6 +4,8 @@
from langchain.agents import agent
from langchain_core.messages import get_buffer_string
from langchain_core.runnables import RunnableConfig
from langchain_groq import ChatGroq
from langchain_ollama import ChatOllama
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import END, START
32 changes: 27 additions & 5 deletions travis/tools.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,37 @@
import os
import uuid
from typing import List

from langchain_community.tools import TavilySearchResults
from langchain_postgres import PGVector
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langchain_core.vectorstores import InMemoryVectorStore
# from langchain_core.vectorstores import InMemoryVectorStore
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document
from dotenv import load_dotenv
load_dotenv()

username = os.getenv("POSTGRES_USER", "sayvai")
password = os.getenv("POSTGRES_PASSWORD", "8056896266")
database = os.getenv("POSTGRES_DB", "user-conversations")
host = os.getenv("POSTGRES_HOST", "postgres")
port = os.getenv("POSTGRES_PORT", "5432")

# Create the connection string
PGVECTOR_URL = f"postgresql+psycopg2://{username}:{password}@{host}:{port}/{database}"

try:
recall_vector_store = PGVector(
connection=PGVECTOR_URL,
embeddings=OpenAIEmbeddings(),
collection_name="document_vectors" # Optional: specify a collection name
)
except Exception as e:
print(f"Failed to connect to vector store: {str(e)}")
raise ValueError(f"Failed to connect to vector store. {str(e)}"
)


def get_user_id(config: RunnableConfig) -> str:
user_id = config["configurable"].get("user_id")
@@ -18,7 +40,7 @@ def get_user_id(config: RunnableConfig) -> str:

return user_id

recall_vector_store = InMemoryVectorStore(OpenAIEmbeddings())
# recall_vector_store = InMemoryVectorStore(OpenAIEmbeddings())

@tool
def save_recall_memory(memory: str, config: RunnableConfig) -> str:
@@ -36,11 +58,11 @@ def search_recall_memories(query: str, config: RunnableConfig) -> List[str]:
"""Search for relevant memories."""
user_id = get_user_id(config)

def _filter_function(doc: Document) -> bool:
return doc.metadata.get("user_id") == user_id
# def _filter_function(doc: Document) -> bool:
# return doc.metadata.get("user_id") == user_id

documents = recall_vector_store.similarity_search(
query, k=3, filter=_filter_function
query, k=3, filter={"user_id": {"$eq": user_id}}
)
return [document.page_content for document in documents]

0 comments on commit 6ee33b2

Please sign in to comment.