LangChain Framework

#!/usr/bin/env python3
"""
Complete LangChain + YugabyteDB Integration Example
This script demonstrates a full end-to-end RAG application.
"""

import os
import getpass
from langchain_yugabytedb import YBEngine, YugabyteDBVectorStore
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.documents import Document
from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

def setup_connection():
    """Set up connection parameters and initialize the vector store."""
    print("Setting up connection to YugabyteDB...")

    # Connection parameters
    YUGABYTEDB_USER = "yugabyte"
    YUGABYTEDB_PASSWORD = ""
    YUGABYTEDB_HOST = "localhost"
    YUGABYTEDB_PORT = "5433"
    YUGABYTEDB_DB = "yugabyte"

    # Table and vector parameters
    TABLE_NAME = "yugabyte_docs_collection"
    VECTOR_SIZE = 1536

    # Create connection string - using psycopg instead of asyncpg for better vector support
    CONNECTION_STRING = (
        f"postgresql+psycopg://{YUGABYTEDB_USER}:{YUGABYTEDB_PASSWORD}@{YUGABYTEDB_HOST}"
        f":{YUGABYTEDB_PORT}/{YUGABYTEDB_DB}"
    )

    # Initialize engine
    engine = YBEngine.from_connection_string(url=CONNECTION_STRING)

    # Initialize embeddings
    print("Initializing OpenAI embeddings...")
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        api_key = getpass.getpass("Enter your OpenAI API Key: ")

    embeddings = OpenAIEmbeddings(api_key=api_key)

    # Initialize vector store table
    print("Creating vector store table...")
    engine.init_vectorstore_table(
        table_name=TABLE_NAME,
        vector_size=VECTOR_SIZE,
    )

    # Create vector store
    vectorstore = YugabyteDBVectorStore.create_sync(
        engine=engine,
        table_name=TABLE_NAME,
        embedding_service=embeddings,
    )

    print("Connection setup complete!")
    return vectorstore, embeddings, api_key

def test_basic_operations(vectorstore):
    """Test basic vector store operations."""
    print("\nTesting basic vector store operations...")

    # Add test documents
    print("Adding test documents...")
    docs = [
        Document(page_content="Apples and oranges are delicious fruits"),
        Document(page_content="Cars and airplanes are modes of transportation"),
        Document(page_content="Trains are efficient for long-distance travel"),
        Document(page_content="YugabyteDB is a distributed SQL database"),
        Document(page_content="Vector databases store embeddings for similarity search"),
        Document(page_content="CDC Logical Replication support, Colocated tables with tablespaces, Auto Analyze service are some features that went GA in 2025.2.0.0."),
    ]

    vectorstore.add_documents(docs)
    print(f"Added {len(docs)} documents to vector store")

    # Test similarity search
    print("\nTesting similarity search...")
    queries = [
        "I'd like to eat some fruit",
        "What can I use to travel long distances?",
        "Tell me about database technology"
    ]

    for query in queries:
        print(f"\nQuery: '{query}'")
        results = vectorstore.similarity_search(query, k=2)
        for i, doc in enumerate(results, 1):
            print(f"  Result {i}: {doc.page_content}")

    return vectorstore

def create_rag_chain(vectorstore, api_key):
    """Create a RAG chain for advanced question answering."""
    print("\nSetting up RAG chain...")

    # Create retriever
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

    # Initialize LLM
    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, api_key=api_key)

    # Define RAG prompt template
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful assistant named 'YugaAI'. Answer questions based on the provided context. If you don't know the answer based on the context, say so."),
        ("human", "Context: {context}\n\nQuestion: {question}")
    ])

    # Build RAG chain
    rag_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    print("RAG chain created successfully!")
    return rag_chain

def test_rag_chain(rag_chain):
    """Test the RAG chain with various questions."""
    print("\nTesting RAG chain...")

    test_questions = [
        "What are some examples of fruits?",
        "What transportation methods were mentioned?",
        "What is YugabyteDB?",
        "How do vector databases work?",
        "Can you tell me all the features which went GA in 2025.2.0.0 release?"
    ]

    for question in test_questions:
        print(f"\nQuestion: {question}")
        try:
            response = rag_chain.invoke(question)
            print(f"YugaAI: {response}")
        except Exception as e:
            print(f"Error: {e}")

def load_yugabyte_docs(vectorstore, embeddings, api_key):
    """Load YugabyteDB documentation and create a more comprehensive vector store."""
    print("\nLoading YugabyteDB documentation...")

    # Load documentation from web
    url = "https://docs.yugabyte.com/stable/releases/ybdb-releases/v2025.2/#v2025.2.0.0"
    loader = WebBaseLoader(url)
    documents = loader.load()

    print(f"Loaded {len(documents)} documents from web")

    # Split documents into chunks
    text_splitter = CharacterTextSplitter(
        separator="\n\n",
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
        is_separator_regex=False,
    )

    chunks = text_splitter.split_documents(documents)
    print(f"Split into {len(chunks)} chunks")

    # Add chunks to vector store
    print("Adding documentation chunks to vector store...")
    vectorstore.add_documents(chunks)
    print(f"Added {len(chunks)} documentation chunks")

    return vectorstore

def main():
    """Main function to run the complete example."""
    print("Starting LangChain + YugabyteDB Complete Example")
    print("=" * 60)

    try:
        # Step 1: Setup
        vectorstore, embeddings, api_key = setup_connection()

        # Step 2: Test basic operations
        vectorstore = test_basic_operations(vectorstore)

        # Step 3: Create RAG chain
        rag_chain = create_rag_chain(vectorstore, api_key)

        # Step 4: Test RAG chain
        test_rag_chain(rag_chain)

        # Step 5: Load additional documentation (optional)
        print("\n" + "=" * 60)
        print("OPTIONAL: Loading YugabyteDB documentation...")
        try:
            vectorstore = load_yugabyte_docs(vectorstore, embeddings, api_key)
            print("Documentation loaded successfully!")
            print("\nTesting with documentation...")
            test_rag_chain(rag_chain)
        except Exception as e:
            print(f"Could not load documentation: {e}")
            print("Continuing with basic example...")

        print("\n" + "=" * 60)
        print("Complete example finished successfully!")
        print("You now have a working RAG application with YugabyteDB!")

    except Exception as e:
        print(f"Error: {e}")
        print("Please check your setup and try again.")

if __name__ == "__main__":
    main()

python langchain_example.py

# Check if YugabyteDB is running
docker ps | grep yugabyte

# Check if vector extension is installed
docker exec -it yugabyte_node01 bin/ysqlsh -h yugabyte01 -c "SELECT * FROM pg_extension WHERE extname = 'vector';"

# Test vector functionality
docker exec -it yugabyte_node01 bin/ysqlsh -h yugabyte01 -c "SELECT '[1,2,3]'::vector;"

# Test Python imports
python -c "from langchain_yugabytedb import YBEngine, YugabyteDBVectorStore; print('Imports successful')"

# Delete specific documents by their IDs
vectorstore.delete(ids=["275823d2-1a47-440d-904b-c07b132fd72b", "another-id-here"])

# Retrieve documents with similarity scores
query = "I'd like a fruit."
docs_with_scores = vectorstore.similarity_search_with_score(query, k=3)

for doc, score in docs_with_scores:
    print(f"Score: {score}")
    print(f"Content: {doc.page_content}\n")

# Create a retriever with custom search parameters
retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 5, "score_threshold": 0.7}
)

# Use in a chain
results = retriever.invoke("What is YugabyteDB?")

import uuid
from langchain_core.messages import SystemMessage, AIMessage, HumanMessage
from langchain_yugabytedb import YugabyteDBChatMessageHistory
import psycopg

# Establish a synchronous connection to the database
conn_info = "dbname=yugabyte user=yugabyte host=localhost port=5433"
sync_connection = psycopg.connect(conn_info)

# Create the table schema (only needs to be done once)
table_name = "chat_history"
YugabyteDBChatMessageHistory.create_tables(sync_connection, table_name)

# Generate a unique session ID
session_id = str(uuid.uuid4())

# Initialize the chat history manager
chat_history = YugabyteDBChatMessageHistory(
    table_name, session_id, sync_connection=sync_connection
)

# Add messages to the chat history
chat_history.add_messages(
    [
        SystemMessage(content="You are a helpful assistant."),
        HumanMessage(content="What is YugabyteDB?"),
        AIMessage(content="YugabyteDB is a distributed SQL database."),
    ]
)

# Retrieve all messages
print(chat_history.messages)

# Add user question to history
chat_history.add_user_message("What features are in v2025.2.0.0?")

# Get conversation context
conversation_context = chat_history.messages

# Use RAG chain to get answer
response = rag_chain.invoke("What features are in v2025.2.0.0?")

# Add AI response to history
chat_history.add_ai_message(response)